Remove test model and use only flow test (the test module is redundant since onnx memory info is saved in redis' INFO report, that can be produced by calling the INFO MODULES command)

Author: alonre24

src/CMakeLists.txt

@@ -1,4 +1,3 @@
-set(CMAKE_BUILD_TYPE Debug)
 if (CMAKE_BUILD_TYPE STREQUAL Debug)
 	SET(DEBUG_SRC "${CMAKE_CURRENT_SOURCE_DIR}/../opt/readies/cetara/diag/gdb.c")
 endif()

src/backends/onnxruntime.c

@@ -17,6 +17,33 @@ OrtAllocator *global_allocator = NULL;
 unsigned long long OnnxMemory = 0;
 unsigned long long OnnxMemoryAccessCounter = 0;
 
+const OrtMemoryInfo *AllocatorInfo(const OrtAllocator *allocator) {
+    (void)allocator;
+    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
+    OrtMemoryInfo *mem_info;
+    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
+        return NULL;
+    }
+    return mem_info;
+}
+
+void *AllocatorAlloc(OrtAllocator *ptr, size_t size) {
+    (void)ptr;
+    void *p = RedisModule_Alloc(size);
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return p;
+}
+
+void AllocatorFree(OrtAllocator *ptr, void *p) {
+    (void)ptr;
+    size_t allocated_size = RedisModule_MallocSize(p);
+    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
+    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
+    return RedisModule_Free(p);
+}
+
 unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; }
 
 unsigned long long RAI_GetMemoryAccessORT() { return OnnxMemoryAccessCounter; }
@@ -269,33 +296,6 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long
     return NULL;
 }
 
-const OrtMemoryInfo *myInfo(const OrtAllocator *allocator) {
-    (void)allocator;
-    const OrtApi *ort = OrtGetApiBase()->GetApi(1);
-    OrtMemoryInfo *mem_info;
-    if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) {
-        return NULL;
-    }
-    return mem_info;
-}
-
-void *myAlloc(OrtAllocator *ptr, size_t size) {
-    (void)ptr;
-    void *p = RedisModule_Alloc(size);
-    size_t allocated_size = RedisModule_MallocSize(p);
-    atomic_fetch_add(&OnnxMemory, allocated_size);
-    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return p;
-}
-
-void myFree(OrtAllocator *ptr, void *p) {
-    (void)ptr;
-    size_t allocated_size = RedisModule_MallocSize(p);
-    atomic_fetch_add(&OnnxMemory, -1 * allocated_size);
-    atomic_fetch_add(&OnnxMemoryAccessCounter, 1);
-    return RedisModule_Free(p);
-}
-
 RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts,
                               const char *modeldef, size_t modellen, RAI_Error *error) {
 
@@ -308,8 +308,8 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo
 
     if (env == NULL) {
         ONNX_API(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env))
-        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, myAlloc, myFree, myInfo,
-                                                  &global_allocator))
+        ONNX_API(ort->CreateCustomDeviceAllocator(ORT_API_VERSION, AllocatorAlloc, AllocatorFree,
+                                                  AllocatorInfo, &global_allocator))
         ONNX_API(ort->RegisterCustomDeviceAllocator(env, global_allocator))
     }
 

src/redisai.c

@@ -23,7 +23,6 @@
 #include <sys/resource.h>
 #include <sys/time.h>
 #include <unistd.h>
-#include "backends/onnxruntime.h"
 
 #include "rmutil/alloc.h"
 #include "rmutil/args.h"
@@ -1093,9 +1092,6 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(DAGRunOpFree, ctx);
     REGISTER_API(DAGFree, ctx);
 
-    // For ORT test module
-    REGISTER_API(LoadDefaultBackend, ctx);
-
     return REDISMODULE_OK;
 }
 
@@ -1112,6 +1108,9 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
                                           RAI_backends.onnx.get_memory_info());
         RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num",
                                           RAI_backends.onnx.get_memory_access_num());
+    } else {
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory", 0);
+        RedisModule_InfoAddFieldULongLong(ctx, "onnxruntime_memory_access_num", 0);
     }
 
     struct rusage self_ru, c_ru;

src/tensor.h

@@ -11,7 +11,7 @@
 #define SRC_TENSOR_H_
 
 #include "config.h"
-//#include "dlpack/dlpack.h"
+#include "dlpack/dlpack.h"
 #include "err.h"
 #include "redismodule.h"
 #include "redisai.h"

src/tensor_struct.h

@@ -1,8 +1,7 @@
 #pragma once
 
 #include "config.h"
-#include "../deps/linux-x64-cpu/dlpack/include/dlpack/dlpack.h"
-//#include "dlpack/dlpack.h"
+#include "dlpack/dlpack.h"
 #include "limits.h"
 
 #define LEN_UNKOWN ULONG_MAX

tests/flow/includes.py

@@ -15,9 +15,6 @@
 import paella
 
 ROOT = os.environ.get("ROOT", None)
-if ROOT is None:
-    sys.stderr.write("ROOT was not defined in the environment.\n")
-    sys.exit(1)
 MAX_ITERATIONS = 2 if os.environ.get("MAX_ITERATIONS") == None else os.environ.get("MAX_ITERATIONS")
 TEST_TF = os.environ.get("TEST_TF") != "0" and os.environ.get("WITH_TF") != "0"
 TEST_TFLITE = os.environ.get("TEST_TFLITE") != "0" and os.environ.get("WITH_TFLITE") != "0"

tests/module/test_data/mul_1.onnx renamed to tests/flow/test_data/mul_1.onnx

@@ -4,33 +4,12 @@
 import redis
 from includes import *
 from RLTest import Env
-from functools import wraps
 
 '''
 python -m RLTest --test tests_onnx.py --module path/to/redisai.so
 '''
 
 
-def load_onnx_test_module(f):
-    @wraps(f)
-    def wrapper(env, *args, **kwargs):
-        goal_dir = os.path.join(os.path.dirname(__file__), "../module/OnnxAllocator.so")
-        TEST_MODULE_PATH = os.path.abspath(goal_dir)
-        con = env.getConnection()
-        modules = con.execute_command("MODULE", "LIST")
-        if b'RAI_onnxAllocator' in [module[1] for module in modules]:
-            return f(env, *args, **kwargs)
-        try:
-            ret = con.execute_command('MODULE', 'LOAD', TEST_MODULE_PATH)
-            env.assertEqual(ret, b'OK')
-        except Exception as e:
-            env.assertFalse(True)
-            env.debugPrint(str(e), force=True)
-            return
-        return f(env, *args, **kwargs)
-    return wrapper
-
-
 def test_onnx_modelrun_mnist(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
@@ -475,17 +454,40 @@ def test_parallelism():
     env.assertEqual(load_time_config["ai_intra_op_parallelism"], "2")
 
 
-@load_onnx_test_module
 def test_onnx_use_custom_allocator(env):
     if not TEST_ONNX:
         env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True)
         return
+    if DEVICE != 'CPU':
+        env.debugPrint("skipping {} since the custom allocator is for CPU".format(sys._getframe().f_code.co_name), force=True)
+        return
 
     con = env.getConnection()
-    ret = con.execute_command("RAI_onnxAllocator.modelSet")
-    env.assertEquals(ret, b'OK')
-
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    model_filename = os.path.join(test_data_path, 'mul_1.onnx')
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                            for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+
+    # Expect using the allocator during model set for allocating the model, its input name and output name.
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertTrue(int(ai_memory_config["ai_onnxruntime_memory"]) > 100)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 3)
+
+    # Expect using the allocator free function when releasing the model.
+    con.execute_command('AI.MODELDEL', 'm{1}')
+    env.assertFalse(con.execute_command('EXISTS', 'm{1}'))
+    ai_memory_config = {k.split(":")[0]: k.split(":")[1]
+                        for k in con.execute_command("INFO MODULES").decode().split("#")[4].split()[1:]}
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory"]), 0)
+    env.assertEqual(int(ai_memory_config["ai_onnxruntime_memory_access_num"]), 6)
+
+    # test allocator in model run op
     model_filename = os.path.join(test_data_path, 'mnist.onnx')
     sample_filename = os.path.join(test_data_path, 'one.raw')
 

tests/flow/tests_onnx.py

@@ -17,11 +17,7 @@ else
 	SHOBJ_LDFLAGS ?= -bundle -undefined dynamic_lookup
 endif
 
-TEST_MODULES = \
-    LLAPI.so \
-    OnnxAllocator.so
-
-SAVE_PATH = @echo "const char* path = \"$(dir $(abspath $(lastword $(MAKEFILE_LIST))))\";" > path.h
+TEST_MODULES = LLAPI.so
 
 .PHONY: all
 
@@ -31,17 +27,12 @@ all: $(TEST_MODULES)
 	$(MAKE) CFLAGS="-m32" LDFLAGS="-melf_i386"
 
 %.o: %.c
-	${SAVE_PATH}
 	$(CC) $(DEBUGFLAGS) -I../../src -I../../deps -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@
 
 LLAPI.so: LLAPI.o DAG_utils.o
 	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
 	chmod +x LLAPI.so
 
-OnnxAllocator.so: OnnxAllocator.o
-	$(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm
-	chmod +x OnnxAllocator.so
-
 .PHONY: clean
 
 clean: