diff --git a/docs/commands.md b/docs/commands.md index db589c2b5..e088e44ee 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1144,7 +1144,7 @@ The **AI.CONFIG** command sets the value of configuration directives at run-time **Redis API** ``` -AI.CONFIG > | > | > +AI.CONFIG > | > | > | > ``` _Arguments_ @@ -1156,6 +1156,7 @@ _Arguments_ * **TORCH**: The PyTorch backend * **ONNX**: ONNXRuntime backend * **MODEL_CHUNK_SIZE**: Sets the size of chunks (in bytes) in which model payloads are split for serialization, replication and `MODELGET`. Default is `511 * 1024 * 1024`. +* **GET**: Retrieve the current value of the `BACKENDSPATH / MODEL_CHUNK_SIZE` configurations. Note that additional information about the module's runtime configuration can be retrieved as part of Redis' info report via `INFO MODULES` command. _Return_ @@ -1190,3 +1191,10 @@ This sets model chunk size to one megabyte (not recommended): redis> AI.CONFIG MODEL_CHUNK_SIZE 1048576 OK ``` + +This returns the current model chunk size configuration: + +``` +redis> AI.CONFIG GET MODEL_CHUNK_SIZE +1048576 +``` \ No newline at end of file diff --git a/src/redisai.c b/src/redisai.c index 1ae2890d1..419619691 100644 --- a/src/redisai.c +++ b/src/redisai.c @@ -939,7 +939,20 @@ int RedisAI_Config_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, i return RedisModule_ReplyWithError(ctx, "ERR MODEL_CHUNK_SIZE: missing chunk size"); } } - + if (!strcasecmp(subcommand, "GET")) { + if (argc > 2) { + const char *config = RedisModule_StringPtrLen(argv[2], NULL); + if (!strcasecmp(config, "BACKENDSPATH")) { + return RedisModule_ReplyWithCString(ctx, Config_GetBackendsPath()); + } else if (!strcasecmp(config, "MODEL_CHUNK_SIZE")) { + return RedisModule_ReplyWithLongLong(ctx, Config_GetModelChunkSize()); + } else { + return RedisModule_ReplyWithNull(ctx); + } + } else { + return RedisModule_WrongArity(ctx); + } + } return RedisModule_ReplyWithError(ctx, "ERR unsupported subcommand"); } diff --git a/tests/flow/tests_commands.py b/tests/flow/tests_commands.py index a7d1f151c..2c98ef471 100644 --- a/tests/flow/tests_commands.py +++ b/tests/flow/tests_commands.py @@ -536,3 +536,79 @@ def run_model_execute_from_llapi(): info = info_to_dict(con.execute_command('AI.INFO', 'm{1}')) env.assertGreaterEqual(info['calls'], 0) env.assertGreaterEqual(num_parallel_clients, info['calls']) + + +def test_ai_config(env): + if not TEST_PT: + env.debugPrint("skipping {} since TEST_PT=0".format(sys._getframe().f_code.co_name), force=True) + return + + conns = env.getOSSMasterNodesConnectionList() + if env.isCluster(): + env.assertEqual(len(conns), env.shardsCount) + + model = load_file_content('pt-minimal.pt') + + for con in conns: + # Get the default configs. + res = con.execute_command('AI.CONFIG', 'GET', 'BACKENDSPATH') + env.assertEqual(res, None) + res = con.execute_command('AI.CONFIG', 'GET', 'MODEL_CHUNK_SIZE') + env.assertEqual(res, 511*1024*1024) + + # Change the default backends path and load backend. + path = f'{ROOT}/install-{DEVICE.lower()}' + con.execute_command('AI.CONFIG', 'BACKENDSPATH', path) + res = con.execute_command('AI.CONFIG', 'GET', 'BACKENDSPATH') + env.assertEqual(res, path.encode()) + be_info = get_info_section(con, "backends_info") + env.assertEqual(len(be_info), 0) # no backends are loaded. + check_error_message(env, con, 'error loading backend', 'AI.CONFIG', 'LOADBACKEND', 'TORCH', ".") + + res = con.execute_command('AI.CONFIG', 'LOADBACKEND', 'TORCH', "backends/redisai_torch/redisai_torch.so") + env.assertEqual(res, b'OK') + be_info = get_info_section(con, "backends_info") + env.assertEqual(len(be_info), 1) # one backend is loaded now - torch. + + # Set the same model twice on some shard - with and without chunks, and assert equality. + con = get_connection(env, '{1}') + chunk_size = len(model) // 3 + model_chunks = [model[i:i + chunk_size] for i in range(0, len(model), chunk_size)] + con.execute_command('AI.MODELSTORE', 'm1{1}', 'TORCH', DEVICE, 'BLOB', model) + con.execute_command('AI.MODELSTORE', 'm2{1}', 'TORCH', DEVICE, 'BLOB', *model_chunks) + model1 = con.execute_command('AI.MODELGET', 'm1{1}', 'BLOB') + model2 = con.execute_command('AI.MODELGET', 'm2{1}', 'BLOB') + env.assertEqual(model1, model2) + + for con in conns: + # Change the default model_chunk_size. + ret = con.execute_command('AI.CONFIG', 'MODEL_CHUNK_SIZE', chunk_size) + env.assertEqual(ret, b'OK') + res = con.execute_command('AI.CONFIG', 'GET', 'MODEL_CHUNK_SIZE') + env.assertEqual(res, chunk_size) + + # Verify that AI.MODELGET returns the model's blob in chunks, with or without the META arg. + con = get_connection(env, '{1}') + model2 = con.execute_command('AI.MODELGET', 'm1{1}', 'BLOB') + env.assertEqual(len(model2), len(model_chunks)) + env.assertTrue(all([el1 == el2 for el1, el2 in zip(model2, model_chunks)])) + + model3 = con.execute_command('AI.MODELGET', 'm1{1}', 'META', 'BLOB')[-1] # Extract the BLOB list from the result + env.assertEqual(len(model3), len(model_chunks)) + env.assertTrue(all([el1 == el2 for el1, el2 in zip(model3, model_chunks)])) + + +def test_ai_config_errors(env): + con = get_connection(env, '{1}') + + check_error_message(env, con, "wrong number of arguments for 'AI.CONFIG' command", 'AI.CONFIG') + check_error_message(env, con, 'unsupported subcommand', 'AI.CONFIG', "bad_subcommand") + check_error_message(env, con, "wrong number of arguments for 'AI.CONFIG' command", 'AI.CONFIG', 'LOADBACKEND') + check_error_message(env, con, 'unsupported backend', 'AI.CONFIG', 'LOADBACKEND', 'bad_backend', "backends/redisai_torch/redisai_torch.so") + check_error_message(env, con, "wrong number of arguments for 'AI.CONFIG' command", 'AI.CONFIG', 'LOADBACKEND', "TORCH") + + check_error_message(env, con, 'BACKENDSPATH: missing path argument', 'AI.CONFIG', 'BACKENDSPATH') + check_error_message(env, con, 'MODEL_CHUNK_SIZE: missing chunk size', 'AI.CONFIG', 'MODEL_CHUNK_SIZE') + + check_error_message(env, con, "wrong number of arguments for 'AI.CONFIG' command", 'AI.CONFIG', 'GET') + env.assertEqual(con.execute_command('AI.CONFIG', 'GET', 'bad_config'), None) diff --git a/tests/flow/tests_pytorch.py b/tests/flow/tests_pytorch.py index 40856660a..f549dcb35 100644 --- a/tests/flow/tests_pytorch.py +++ b/tests/flow/tests_pytorch.py @@ -9,37 +9,6 @@ ''' -def test_pytorch_chunked_modelstore(env): - if not TEST_PT: - env.debugPrint("skipping {} since TEST_PT=0".format(sys._getframe().f_code.co_name), force=True) - return - - con = get_connection(env, '{1}') - model = load_file_content('pt-minimal.pt') - - chunk_size = len(model) // 3 - - model_chunks = [model[i:i + chunk_size] for i in range(0, len(model), chunk_size)] - - ret = con.execute_command('AI.MODELSTORE', 'm1{1}', 'TORCH', DEVICE, 'BLOB', model) - ret = con.execute_command('AI.MODELSTORE', 'm2{1}', 'TORCH', DEVICE, 'BLOB', *model_chunks) - - model1 = con.execute_command('AI.MODELGET', 'm1{1}', 'BLOB') - model2 = con.execute_command('AI.MODELGET', 'm2{1}', 'BLOB') - - env.assertEqual(model1, model2) - - ret = con.execute_command('AI.CONFIG', 'MODEL_CHUNK_SIZE', chunk_size) - - model2 = con.execute_command('AI.MODELGET', 'm2{1}', 'BLOB') - env.assertEqual(len(model2), len(model_chunks)) - env.assertTrue(all([el1 == el2 for el1, el2 in zip(model2, model_chunks)])) - - model3 = con.execute_command('AI.MODELGET', 'm2{1}', 'META', 'BLOB')[-1] # Extract the BLOB list from the result - env.assertEqual(len(model3), len(model_chunks)) - env.assertTrue(all([el1 == el2 for el1, el2 in zip(model3, model_chunks)])) - - def test_pytorch_modelrun(env): if not TEST_PT: env.debugPrint("skipping {} since TEST_PT=0".format(sys._getframe().f_code.co_name), force=True)