diff --git a/CMakeLists.txt b/CMakeLists.txt index 632e5c8d6..88510e119 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -160,6 +160,7 @@ ENDIF() IF(BUILD_ORT) FIND_LIBRARY(ORT_LIBRARIES NAMES onnxruntime PATHS ${depsAbs}/onnxruntime/lib) + ADD_SUBDIRECTORY(src/backends/onnx_allocator) MESSAGE(STATUS "Found ONNXRuntime Libraries: \"${ORT_LIBRARIES}\")") IF (NOT ORT_LIBRARIES) MESSAGE(FATAL_ERROR "Could not find ONNXRuntime") @@ -293,6 +294,7 @@ ENDIF() IF(BUILD_ORT) ADD_LIBRARY(redisai_onnxruntime SHARED $) + TARGET_LINK_LIBRARIES(redisai_onnxruntime onnx_allocator ${ORT_LIBRARIES}) TARGET_LINK_LIBRARIES(redisai_onnxruntime ${ORT_LIBRARIES}) SET_TARGET_PROPERTIES(redisai_onnxruntime PROPERTIES PREFIX "") SET_TARGET_PROPERTIES(redisai_onnxruntime PROPERTIES SUFFIX ".so") diff --git a/src/backends/backends.c b/src/backends/backends.c index 02a17bcf8..59ec5a277 100644 --- a/src/backends/backends.c +++ b/src/backends/backends.c @@ -50,6 +50,8 @@ int RAI_ExportFunc(const char *func_name, void **targetFuncPtr) { *targetFuncPtr = Config_GetModelExecutionTimeout; } else if (strcmp("GetThreadsCount", func_name) == 0) { *targetFuncPtr = BGWorker_GetThreadsCount; + } else if (strcmp("GetBackendMemoryLimit", func_name) == 0) { + *targetFuncPtr = Config_GetBackendMemoryLimit; // Export RedisAI low level API functions. } else if (strcmp("RedisAI_InitError", func_name) == 0) { diff --git a/src/backends/backends_api.h b/src/backends/backends_api.h index 7d7919f27..13d2feeb5 100644 --- a/src/backends/backends_api.h +++ b/src/backends/backends_api.h @@ -37,12 +37,20 @@ BACKENDS_API uintptr_t (*RedisAI_GetThreadsCount)(void); BACKENDS_API long long (*RedisAI_GetNumThreadsPerQueue)(void); /** - * @return The maximal number of milliseconds that a model run session should run + * @return The maximum number of milliseconds that a model run session should run * before it is terminated forcefully (load time config). - * Currently supported only fo onnxruntime backend. + * Currently supported only for onnxruntime backend. */ BACKENDS_API long long (*RedisAI_GetModelExecutionTimeout)(void); +/** + * @return The maximum number of memory (in MB) that a backend can consume + * for creating and running inference sessions. When memory limit is exceeded, operation + * is not permitted and an error is returned. + * Currently supported only for onnxruntime backend. + */ +BACKENDS_API long long (*RedisAI_GetMemoryLimit)(void); + /** * The following functions are part of RedisAI low level API (the full low level * API is defined in redisai.h). For every function below named "RedisAI_X", its diff --git a/src/backends/onnx_allocator/CMakeLists.txt b/src/backends/onnx_allocator/CMakeLists.txt new file mode 100644 index 000000000..89503fe4d --- /dev/null +++ b/src/backends/onnx_allocator/CMakeLists.txt @@ -0,0 +1,3 @@ +add_library(onnx_allocator STATIC onnx_allocator.cpp) +target_link_libraries(onnx_allocator "${ONNX_LIBRARIES}") +set_property(TARGET onnx_allocator PROPERTY CXX_STANDARD 14) \ No newline at end of file diff --git a/src/backends/onnx_allocator/onnx_allocator.cpp b/src/backends/onnx_allocator/onnx_allocator.cpp new file mode 100644 index 000000000..432db4893 --- /dev/null +++ b/src/backends/onnx_allocator/onnx_allocator.cpp @@ -0,0 +1,114 @@ +#include "onnx_allocator.h" +#include "../onnxruntime.h" +#include "onnxruntime_cxx_api.h" +#include + +struct RAIOrtAllocator : OrtAllocator { + RAIOrtAllocator(); + ~RAIOrtAllocator(); + RAIOrtAllocator(const RAIOrtAllocator&) = delete; + RAIOrtAllocator& operator=(const RAIOrtAllocator&) = delete; + + void* Alloc(size_t size); + void Free(void* p); + const OrtMemoryInfo* Info() const; + unsigned long long NumAllocatorAccess() const; + unsigned long long MemoryInUse() const; + void SetMemoryLimit(unsigned long long max_memory); + static RAIOrtAllocator *GetInstance(); + +private: + std::atomic memory_inuse{0}; + std::atomic num_allocator_access{0}; + unsigned long long memory_limit = 0; + OrtMemoryInfo* cpu_memory_info; + static RAIOrtAllocator* allocator_instance; +}; + +RAIOrtAllocator* RAIOrtAllocator::allocator_instance = nullptr; + +RAIOrtAllocator::RAIOrtAllocator() { + OrtAllocator::version = ORT_API_VERSION; + OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast(this_)->Alloc(size); }; + OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast(this_)->Free(p); }; + OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast(this_)->Info(); }; + Ort::ThrowOnError(Ort::GetApi().CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &cpu_memory_info)); + RAIOrtAllocator::allocator_instance = this; +} + +RAIOrtAllocator::~RAIOrtAllocator() { + Ort::GetApi().ReleaseMemoryInfo(cpu_memory_info); +} + +void* RAIOrtAllocator::Alloc(size_t size) { + // Allocate an additional 63 bytes to ensure that we can return an address which is + // 64-byte aligned, and an additional space in the size of a pointer to store + // the address that RedisModule_Alloc returns. + int offset = 63 + sizeof(void *); + void *allocated_address = (void *)RedisModule_Alloc(size + offset); + size_t allocated_size = RedisModule_MallocSize(allocated_address); + // Update the total number of bytes that onnx is using and the number of accesses + // that onnx made to the allocator. + size_t cur_memory = memory_inuse.load(); + if (memory_limit && cur_memory + allocated_size > memory_limit) { + RedisModule_Free(allocated_address); + throw Ort::Exception("Onnxruntime memory limit exceeded, memory allocation failed.", ORT_RUNTIME_EXCEPTION); + } + memory_inuse.fetch_add(allocated_size); + num_allocator_access.fetch_add(1); + // This operation guarantees that "aligned_address" is the closest 64-aligned address to ("allocated_address"+size_t). + void **aligned_address = (void **)(((size_t)(allocated_address) + offset) & (~63)); + // This stores the address "allocated_address" right before "aligned_address" (so we can retrieve it when we free). + aligned_address[-1] = allocated_address; + return aligned_address; +} + +void RAIOrtAllocator::Free(void* p) { + if (p == nullptr) { + return; + } + // Retrieve the address that we originally received from RedisModule_Alloc + // (this is the address that we need to sent to RedisModule_Free). + void *allocated_address = ((void **)p)[-1]; + size_t allocated_size = RedisModule_MallocSize(allocated_address); + // Update the total number of bytes that onnx is using and the number of accesses + // that onnx made to the allocator. + memory_inuse.fetch_sub(allocated_size); + num_allocator_access.fetch_add(1); + RedisModule_Free(allocated_address); +} + +const OrtMemoryInfo* RAIOrtAllocator::Info() const { + return cpu_memory_info; +} + +unsigned long long RAIOrtAllocator::NumAllocatorAccess() const { + return num_allocator_access.load(); +} + +unsigned long long RAIOrtAllocator::MemoryInUse() const { + return memory_inuse.load(); +} + +void RAIOrtAllocator::SetMemoryLimit(unsigned long long max_memory) { + // max_memory is given in MB + memory_limit = 1000000*max_memory; +} + +RAIOrtAllocator *RAIOrtAllocator::GetInstance() { + return RAIOrtAllocator::allocator_instance; +} + +OrtAllocator *CreateCustomAllocator(unsigned long long max_memory) { + auto *allocator = new RAIOrtAllocator(); + allocator->SetMemoryLimit(max_memory); + return allocator; +} + +unsigned long long RAI_GetMemoryInfoORT() { + return RAIOrtAllocator::GetInstance()->MemoryInUse(); +} + +unsigned long long RAI_GetMemoryAccessORT() { + return RAIOrtAllocator::GetInstance()->NumAllocatorAccess(); +} diff --git a/src/backends/onnx_allocator/onnx_allocator.h b/src/backends/onnx_allocator/onnx_allocator.h new file mode 100644 index 000000000..f7db13af1 --- /dev/null +++ b/src/backends/onnx_allocator/onnx_allocator.h @@ -0,0 +1,17 @@ +#pragma once + +#include "onnxruntime_c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +OrtAllocator *CreateCustomAllocator(unsigned long long max_memory); + +unsigned long long RAI_GetMemoryInfoORT(); + +unsigned long long RAI_GetMemoryAccessORT(); + +#ifdef __cplusplus +} +#endif diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c index ed3a41efc..403e7e2b0 100644 --- a/src/backends/onnxruntime.c +++ b/src/backends/onnxruntime.c @@ -6,6 +6,7 @@ #include "util/arr.h" #include "backends/onnxruntime.h" #include "redis_ai_objects/tensor.h" +#include "onnx_allocator/onnx_allocator.h" #include "onnxruntime_c_api.h" #include "backends_api.h" @@ -21,63 +22,7 @@ OrtEnv *env = NULL; // For model that run on GPU, onnx will not use the custom allocator (redis allocator), but // the onnx allocator for GPU. But for the auxiliary allocations of the input and output names, // we will use the custom global allocator for models that run on GPU as well. -OrtMemoryInfo *mem_info = NULL; 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); - if (mem_info != NULL) { - return mem_info; - } - if (ort->CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &mem_info) != NULL) { - return NULL; - } - return mem_info; -} - -// Allocate address with 64-byte alignment to cope with onnx optimizations. -void *AllocatorAlloc(OrtAllocator *ptr, size_t size) { - - (void)ptr; - // Allocate an additional 63 bytes to ensure that we can return an address which is - // 64-byte aligned, and an additional space in the size of a pointer to store - // the address that RedisModule_Alloc returns. - int offset = 63 + sizeof(void *); - void *allocated_address = (void *)RedisModule_Alloc(size + offset); - size_t allocated_size = RedisModule_MallocSize(allocated_address); - // Update the total number of bytes that onnx is using and the number of accesses - // that onnx made to the allocator. - atomic_fetch_add(&OnnxMemory, allocated_size); - atomic_fetch_add(&OnnxMemoryAccessCounter, 1); - // This operation guarantees that p2 is the closest 64-aligned address to (p1+size_t). - void **aligned_address = (void **)(((size_t)(allocated_address) + offset) & (~63)); - // This stores the address p1 right before p2 (so we can retrieve it when we free). - aligned_address[-1] = allocated_address; - return aligned_address; -} - -void AllocatorFree(OrtAllocator *ptr, void *aligned_address) { - (void)ptr; - if (aligned_address == NULL) { - return; - } - // Retrieve the address that we originally received from RedisModule_Alloc - // (this is the address that we need to sent to RedisModule_Free). - void *allocated_address = ((void **)aligned_address)[-1]; - size_t allocated_size = RedisModule_MallocSize(allocated_address); - // Update the total number of bytes that onnx is using and the number of accesses - // that onnx made to the allocator. - atomic_fetch_sub(&OnnxMemory, allocated_size); - atomic_fetch_add(&OnnxMemoryAccessCounter, 1); - return RedisModule_Free(allocated_address); -} - -unsigned long long RAI_GetMemoryInfoORT() { return OnnxMemory; } - -unsigned long long RAI_GetMemoryAccessORT() { return OnnxMemoryAccessCounter; } int RAI_InitBackendORT(int (*get_api_fn)(const char *, void **)) { // Export redis callbacks. @@ -95,6 +40,7 @@ int RAI_InitBackendORT(int (*get_api_fn)(const char *, void **)) { get_api_fn("GetThreadId", ((void **)&RedisAI_GetThreadId)); get_api_fn("GetNumThreadsPerQueue", ((void **)&RedisAI_GetNumThreadsPerQueue)); get_api_fn("GetModelExecutionTimeout", ((void **)&RedisAI_GetModelExecutionTimeout)); + get_api_fn("GetBackendMemoryLimit", ((void **)&RedisAI_GetMemoryLimit)); get_api_fn("GetThreadsCount", ((void **)&RedisAI_GetThreadsCount)); // Create a global array of onnx runSessions, with an entry for every working thread. @@ -389,8 +335,9 @@ RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_Mo // allocating buffers when creating and running models that run on CPU, and for allocations of // models inputs and outputs names (for both models that run on CPU and GPU) if (env == NULL) { - ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "test", &env)) - ONNX_VALIDATE_STATUS(ort->GetAllocatorWithDefaultOptions(&global_allocator)); + ONNX_VALIDATE_STATUS(ort->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "RedisAI", &env)) + global_allocator = CreateCustomAllocator(RedisAI_GetMemoryLimit()); + ONNX_VALIDATE_STATUS(ort->RegisterAllocator(env, global_allocator)) } ONNX_VALIDATE_STATUS(ort->CreateSessionOptions(&session_options)) diff --git a/src/backends/onnxruntime.h b/src/backends/onnxruntime.h index d165af32c..36db1b380 100644 --- a/src/backends/onnxruntime.h +++ b/src/backends/onnxruntime.h @@ -5,10 +5,6 @@ #include "redis_ai_objects/model.h" #include "execution/execution_contexts/execution_ctx.h" -unsigned long long RAI_GetMemoryInfoORT(void); - -unsigned long long RAI_GetMemoryAccessORT(void); - int RAI_InitBackendORT(int (*get_api_fn)(const char *, void **)); RAI_Model *RAI_ModelCreateORT(RAI_Backend backend, const char *devicestr, RAI_ModelOpts opts, diff --git a/src/config/config.c b/src/config/config.c index e51d43692..0a4e6a4e3 100644 --- a/src/config/config.c +++ b/src/config/config.c @@ -16,6 +16,8 @@ long long ThreadPoolSizePerQueue = 1; // Number of working threads for devi long long ModelExecutionTimeout = 5000; // The maximum time in milliseconds // before killing onnx run session. +long long BackendMemoryLimit = 0; // The maximum amount of memory in MB + // that backend is allowed to consume. static int _Config_LoadTimeParamParse(RedisModuleCtx *ctx, const char *key, const char *val, RedisModuleString *rsval) { @@ -56,6 +58,11 @@ static int _Config_LoadTimeParamParse(RedisModuleCtx *ctx, const char *key, cons if (ret == REDISMODULE_OK) { RedisModule_Log(ctx, "notice", "%s: %s", REDISAI_INFOMSG_MODEL_EXECUTION_TIMEOUT, val); } + } else if (strcasecmp((key), "BACKEND_MEMORY_LIMIT") == 0) { + ret = Config_SetBackendMemoryLimit(rsval); + if (ret == REDISMODULE_OK) { + RedisModule_Log(ctx, "notice", "%s: %s", REDISAI_INFOMSG_BACKEND_MEMORY_LIMIT, val); + } } else if (strcasecmp((key), "BACKENDSPATH") == 0) { // already taken care of } else { @@ -74,6 +81,8 @@ long long Config_GetNumThreadsPerQueue() { return ThreadPoolSizePerQueue; } long long Config_GetModelExecutionTimeout() { return ModelExecutionTimeout; } +long long Config_GetBackendMemoryLimit() { return BackendMemoryLimit; } + char *Config_GetBackendsPath() { return BackendsPath; } int Config_LoadBackend(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { @@ -160,6 +169,16 @@ int Config_SetModelExecutionTimeout(RedisModuleString *timeout) { return REDISMODULE_OK; } +int Config_SetBackendMemoryLimit(RedisModuleString *memory_limit) { + long long val; + int result = RedisModule_StringToLongLong(memory_limit, &val); + if (result != REDISMODULE_OK || val <= 0) { + return REDISMODULE_ERR; + } + BackendMemoryLimit = val; + return REDISMODULE_OK; +} + int Config_SetLoadTimeParams(RedisModuleCtx *ctx, RedisModuleString *const *argv, int argc) { if (argc > 0 && argc % 2 != 0) { RedisModule_Log(ctx, "warning", diff --git a/src/config/config.h b/src/config/config.h index 724cc509a..d6798f2f3 100644 --- a/src/config/config.h +++ b/src/config/config.h @@ -26,6 +26,7 @@ typedef enum { RAI_DEVICE_CPU = 0, RAI_DEVICE_GPU = 1 } RAI_Device; #define REDISAI_INFOMSG_INTER_OP_PARALLELISM "Setting INTER_OP_PARALLELISM parameter to" #define REDISAI_INFOMSG_MODEL_CHUNK_SIZE "Setting MODEL_CHUNK_SIZE parameter to" #define REDISAI_INFOMSG_MODEL_EXECUTION_TIMEOUT "Setting MODEL_EXECUTION_TIMEOUT parameter to" +#define REDISAI_INFOMSG_BACKEND_MEMORY_LIMIT "Setting BACKEND_MEMORY_LIMIT parameter to" /** * Get number of threads used for parallelism between independent operations, by @@ -56,6 +57,13 @@ long long Config_GetNumThreadsPerQueue(void); */ long long Config_GetModelExecutionTimeout(void); +/** + * @return Memory limit in MB for backend. This is the maximum amount of memory + * that can be consumed by the backend for creating and running sessions. + * Currently supported only for onnxruntime backend. + */ +long long Config_GetBackendMemoryLimit(void); + /** * @return Returns the backends path string. */ @@ -113,11 +121,19 @@ int Config_SetModelChunkSize(RedisModuleString *chunk_size_string); /** * Set the maximum time in ms that onnx backend allow running a model. - * @param onnx_max_runtime - string containing the max runtime (in ms) + * @param timeout - string containing the max runtime (in ms) * @return REDISMODULE_OK on success, or REDISMODULE_ERR if failed */ int Config_SetModelExecutionTimeout(RedisModuleString *timeout); +/** + * Set the memory limit in MB for backends allocations. + * @param memory_limit - maximum memory consumption by backend. If values is zero, + * there will be no enforcement of any memory limit. + * @return REDISMODULE_OK on success, or REDISMODULE_ERR if failed + */ +int Config_SetBackendMemoryLimit(RedisModuleString *memory_limit); + /** * Load time configuration parser * @param ctx Context in which Redis modules operate diff --git a/src/redisai.c b/src/redisai.c index ee71483ed..e3639aa66 100644 --- a/src/redisai.c +++ b/src/redisai.c @@ -1167,6 +1167,7 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) { Config_GetBackendsIntraOpParallelism()); RedisModule_InfoAddFieldLongLong(ctx, "model_execution_timeout", Config_GetModelExecutionTimeout()); + RedisModule_InfoAddFieldLongLong(ctx, "backend_memory_limit", Config_GetBackendMemoryLimit()); _moduleInfo_getBackendsInfo(ctx); struct rusage self_ru, c_ru; diff --git a/src/util/arr.h b/src/util/arr.h index fd4427189..788cdf510 100644 --- a/src/util/arr.h +++ b/src/util/arr.h @@ -58,7 +58,7 @@ 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_alloc_fn(sizeof(array_hdr_t) + cap * elem_sz); + array_hdr_t *hdr = (array_hdr_t *)array_alloc_fn(sizeof(array_hdr_t) + cap * elem_sz); hdr->on_stack = false; hdr->cap = cap; hdr->elem_sz = elem_sz; @@ -116,7 +116,7 @@ static inline array_t array_ensure_cap(array_t arr, uint32_t cap) { } if (cap > hdr->cap) { hdr->cap = MAX(hdr->cap * 2, cap); - hdr = array_realloc_fn(hdr, array_sizeof(hdr)); + hdr = (array_hdr_t *)array_realloc_fn(hdr, array_sizeof(hdr)); } return (array_t)hdr->buf; } @@ -150,7 +150,7 @@ static inline void *array_trimm(array_t arr, uint32_t len, uint32_t cap) { if (cap != -1) { arr_hdr->cap = cap; if (!arr_hdr->on_stack) - arr_hdr = array_realloc_fn(arr_hdr, array_sizeof(arr_hdr)); + arr_hdr = (array_hdr_t *)array_realloc_fn(arr_hdr, array_sizeof(arr_hdr)); } return arr_hdr->buf; } diff --git a/tests/flow/test_data/inception-v2-9.onnx b/tests/flow/test_data/inception-v2-9.onnx new file mode 100644 index 000000000..21aafc68a --- /dev/null +++ b/tests/flow/test_data/inception-v2-9.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37ed4aa93f817fd0d1cc19bda045ebd70f18d0887a5f51e66ecd0ed8fd5a65c2 +size 45042799 diff --git a/tests/flow/tests_common.py b/tests/flow/tests_common.py index 8ca5d1c05..9e0579e2b 100644 --- a/tests/flow/tests_common.py +++ b/tests/flow/tests_common.py @@ -303,7 +303,8 @@ def test_info_command(env): env.assertEqual(list(git.keys()), ['ai_git_sha']) load_time_configs = get_info_section(con, 'load_time_configs') env.assertEqual(list(load_time_configs.keys()), ['ai_threads_per_queue', 'ai_inter_op_parallelism', - 'ai_intra_op_parallelism', 'ai_model_execution_timeout']) + 'ai_intra_op_parallelism', 'ai_model_execution_timeout', + 'ai_backend_memory_limit']) # minimum cpu properties cpu = get_info_section(con, 'cpu') env.assertTrue('ai_self_used_cpu_sys' in cpu.keys()) diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py index 7500c9baa..f79bb3e9b 100644 --- a/tests/flow/tests_onnx.py +++ b/tests/flow/tests_onnx.py @@ -364,6 +364,130 @@ def test_parallelism(): env.assertEqual(load_time_config["ai_intra_op_parallelism"], "2") +class TestOnnxCustomAllocator: + def __init__(self): + self.env = Env() + if not TEST_ONNX: + self.env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) + return + self.allocator_access_counter = 0 + + def test_1_cpu_allocator(self): + con = get_connection(self.env, '{1}') + model_pb = load_file_content('mul_1.onnx') + + # Expect using the allocator during model set for allocating the model, its input name and output name: + # overall 3 allocations. The model raw size is 24B ,and the names are 2B each. In practice we allocate + # more than 28B as Redis allocator will use additional memory for its internal management and for the + # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations + # (hence will not use additional memory). + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'ONNX', 'CPU', 'BLOB', model_pb) + self.env.assertEqual(ret, b'OK') + self.allocator_access_counter += 3 + backends_info = get_info_section(con, 'backends_info') + + # Expect using at least 24+63+(size of an address) + 2*(2+63+(size of an address)) (=241) bytes. + model_allocation_bytes_used = int(backends_info["ai_onnxruntime_memory"]) + self.env.assertTrue(model_allocation_bytes_used >= 241) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + con.execute_command('AI.TENSORSET', 'a_mul{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0) + + # Running the model should access the allocator 6 times: allocating+freeing input+output names, + # and allocating+freeing the output as OrtValue. Overall, there should be no change in the memory consumption. + con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'a_mul{1}', 'OUTPUTS', 1, 'b{1}') + self.allocator_access_counter += 6 + values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES') + self.env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36']) + backends_info = get_info_section(con, 'backends_info') + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory"]), model_allocation_bytes_used) + + # Expect using the allocator free function 3 times: when releasing the model, input name and output name. + con.execute_command('AI.MODELDEL', 'm{1}') + self.allocator_access_counter += 3 + self.env.assertFalse(con.execute_command('EXISTS', 'm{1}')) + backends_info = get_info_section(con, 'backends_info') + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory"]), 0) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + + def test_2_with_gpu(self): + if DEVICE == 'CPU': + self.env.debugPrint("skipping {} since this test if for GPU only".format(sys._getframe().f_code.co_name), force=True) + return + con = get_connection(self.env, '{1}') + model_pb = load_file_content('mul_1.onnx') + + # for GPU, expect using the allocator only for allocating input and output names (not the model itself). + ret = con.execute_command('AI.MODELSTORE', 'm_gpu{1}', 'ONNX', DEVICE, 'BLOB', model_pb) + self.env.assertEqual(ret, b'OK') + self.allocator_access_counter += 2 + + # Expect using at least 2*(2+63+(size of an address))(=146) bytes by redis allocator, but no more than 240, + # as the model weights shouldn't be allocated by the allocator. + backends_info = get_info_section(con, 'backends_info') + model_allocation_bytes_used = int(backends_info["ai_onnxruntime_memory"]) + self.env.assertTrue(model_allocation_bytes_used > 146) + self.env.assertTrue(model_allocation_bytes_used < 241) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + + # Make sure that allocator is not used for running and freeing the GPU model, except for + # the input and output names allocations (and deallocations). + con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 3, 2, 'VALUES', 1.0, 2.0, 3.0, 4.0, 5.0, 6.0) + con.execute_command('AI.MODELEXECUTE', 'm_gpu{1}', 'INPUTS', 1, 'a{1}', 'OUTPUTS', 1, 'b{1}') + self.allocator_access_counter += 4 + values = con.execute_command('AI.TENSORGET', 'b{1}', 'VALUES') + self.env.assertEqual(values, [b'1', b'4', b'9', b'16', b'25', b'36']) + + # Expect that memory usage didn't change, and for another 4 accesses to the allocator (input and output names + # allocation and free) + backends_info = get_info_section(con, 'backends_info') + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory"]), model_allocation_bytes_used) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + + # Expect only 2 more accesses in delete - for deallocating input and output names + con.execute_command('AI.MODELDEL', 'm_gpu{1}') + self.allocator_access_counter += 2 + self.env.assertFalse(con.execute_command('EXISTS', 'm_gpu{1}')) + backends_info = get_info_section(con, 'backends_info') + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory"]), 0) + self.env.assertEqual(int(backends_info["ai_onnxruntime_memory_access_num"]), self.allocator_access_counter) + + def test_3_memory_limit(self): + self.env = Env(moduleArgs='THREADS_PER_QUEUE 8 BACKEND_MEMORY_LIMIT 1') + self.allocator_access_counter = 0 + con = get_connection(self.env, '{1}') + + # Try to allocate a model whose size exceeds the memory limit + inception_pb = load_file_content('inception-v2-9.onnx') + check_error_message(self.env, con, "Exception during initialization: Onnxruntime memory limit exceeded," + " memory allocation failed.", + 'AI.MODELSTORE', 'inception{1}', 'ONNX', 'CPU', 'BLOB', inception_pb) + + mnist_pb = load_file_content('mnist.onnx') + sample_raw = load_file_content('one.raw') + + # Create 25 different sessions of mnist model, the size of each session in onnx is ~31KB, overall ~770KB + for i in range(25): + ret = con.execute_command('AI.MODELSTORE', 'mnist_'+str(i)+'{1}', 'ONNX', 'CPU', 'BLOB', mnist_pb) + self.env.assertEqual(ret, b'OK') + con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw) + + # As onnx memory consumption is about 0.77MB at this point, and executing mnist session requires an additional + # 500KB of memory, we are expected to exceed the memory limit here in some operation. Note that the exact + # memory consumption here changes whether we are using libc allocator or jemalloc (jemalloc will be greater) + check_error_message(self.env, con, "Onnxruntime memory limit exceeded, memory allocation failed.", + 'AI.MODELEXECUTE', 'mnist_0{1}', 'INPUTS', 1, 'a{1}', 'OUTPUTS', 1, 'b{1}', + error_msg_is_substr=True) + + def run_parallel_onnx_sessions(con): + check_error_message(self.env, con, "Onnxruntime memory limit exceeded, memory allocation failed.", + 'AI.MODELEXECUTE', 'mnist_0{1}', 'INPUTS', 1, 'a{1}', 'OUTPUTS', 1, 'b{1}', + error_msg_is_substr=True) + + # We run sessions in parallel, all of them should fail. Note that here. + run_test_multiproc(self.env, '{1}', 50, run_parallel_onnx_sessions) + + class TestOnnxKillSwitch: def __init__(self):