Add possibility to provide a model in chunks (#338)

* Add possibility to provide a model in chunks

* Add test on chunked modelset

Author: lantiga

README.md

@@ -47,7 +47,7 @@ Note that Redis config is located at `/usr/local/etc/redis/redis.conf` which can
 
 On the client, set the model
 ```sh
-redis-cli -x AI.MODELSET foo TF CPU INPUTS a b OUTPUTS c < test/test_data/graph.pb
+redis-cli -x AI.MODELSET foo TF CPU INPUTS a b OUTPUTS c BLOB < test/test_data/graph.pb
 ```
 
 Then create the input tensors, run the computation graph and get the output tensor (see `load_model.sh`). Note the signatures:

docs/commands.md

@@ -123,7 +123,7 @@ AI.TENSORGET foo META VALUES
 Set a model.
 
 ```sql
-AI.MODELSET model_key backend device [TAG tag] [BATCHSIZE n [MINBATCHSIZE m]] [INPUTS name1 name2 ... OUTPUTS name1 name2 ...] model_blob
+AI.MODELSET model_key backend device [TAG tag] [BATCHSIZE n [MINBATCHSIZE m]] [INPUTS name1 name2 ... OUTPUTS name1 name2 ...] BLOB model_blob
 ```
 
 * model_key - Key for storing the model
@@ -142,28 +142,29 @@ AI.MODELSET model_key backend device [TAG tag] [BATCHSIZE n [MINBATCHSIZE m]] [I
                    Default is 0 (no minimum batch size).
 * INPUTS name1 name2 ... - Name of the nodes in the provided graph corresponding to inputs [`TF` backend only]
 * OUTPUTS name1 name2 ... - Name of the nodes in the provided graph corresponding to outputs [`TF` backend only]
-* model_blob - Binary buffer containing the model protobuf saved from a supported backend
+* BLOB model_blob - Binary buffer containing the model protobuf saved from a supported backend. Since Redis supports strings
+                    up to 512MB, blobs for very large models need to be chunked, e.g. `BLOB chunk1 chunk2 ...`.
 
 ### MODELSET Example
 
 ```sql
-AI.MODELSET resnet18 TORCH GPU < foo.pt
+AI.MODELSET resnet18 TORCH GPU BLOB < foo.pt
 ```
 
 ```sql
-AI.MODELSET resnet18 TF CPU INPUTS in1 OUTPUTS linear4 < foo.pb
+AI.MODELSET resnet18 TF CPU INPUTS in1 OUTPUTS linear4 BLOB < foo.pb
 ```
 
 ```sql
-AI.MODELSET mnist_net ONNX CPU TAG mnist:lenet:v0.1 < mnist.onnx
+AI.MODELSET mnist_net ONNX CPU TAG mnist:lenet:v0.1 BLOB < mnist.onnx
 ```
 
 ```sql
-AI.MODELSET mnist_net ONNX CPU BATCHSIZE 10 < mnist.onnx
+AI.MODELSET mnist_net ONNX CPU BATCHSIZE 10 BLOB < mnist.onnx
 ```
 
 ```sql
-AI.MODELSET resnet18 TF CPU BATCHSIZE 10 MINBATCHSIZE 6 INPUTS in1 OUTPUTS linear4 < foo.pb
+AI.MODELSET resnet18 TF CPU BATCHSIZE 10 MINBATCHSIZE 6 INPUTS in1 OUTPUTS linear4 BLOB < foo.pb
 ```
 
 ## AI.MODELGET
@@ -284,13 +285,13 @@ AI._MODELSCAN
 Set a script.
 
 ```sql
-AI.SCRIPTSET script_key device [TAG tag] script_source
+AI.SCRIPTSET script_key device [TAG tag] SOURCE script_source
 ```
 
 * script_key - Key for storing the script
 * device - The device where the script will execute
 * TAG tag - Optional string tagging the script, such as a version number or other identifier
-* script_source - A string containing [TorchScript](https://pytorch.org/docs/stable/jit.html) source code
+* SOURCE script_source - A string containing [TorchScript](https://pytorch.org/docs/stable/jit.html) source code
 
 ### SCRIPTSET Example
 
@@ -302,11 +303,11 @@ def addtwo(a, b):
 ```
 
 ```sql
-AI.SCRIPTSET addscript GPU < addtwo.txt
+AI.SCRIPTSET addscript GPU SOURCE < addtwo.txt
 ```
 
 ```sql
-AI.SCRIPTSET addscript GPU TAG myscript:v0.1 < addtwo.txt
+AI.SCRIPTSET addscript GPU TAG myscript:v0.1 SOURCE < addtwo.txt
 ```
 
 ## AI.SCRIPTGET

docs/index.md

@@ -24,7 +24,7 @@ docker run -p 6379:6379 -it --rm redisai/redisai
 On the client, load a backend (TF, TORCH or ONNX), and set the model
 ```sh
 redis-cli AI.CONFIG LOADBACKEND TF install/backends/redisai_tensorflow/redisai_tensorflow.so
-redis-cli -x AI.MODELSET foo TF CPU INPUTS a b OUTPUTS c < test/test_data/graph.pb
+redis-cli -x AI.MODELSET foo TF CPU INPUTS a b OUTPUTS c BLOB < test/test_data/graph.pb
 ```
 
 Then create the input tensors, run the computation graph and get the output tensor (see `load_model.sh`). Note the signatures:

src/redisai.c

@@ -86,7 +86,7 @@ int RedisAI_TensorGet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
 }
 
 /**
-* AI.MODELSET model_key backend device [TAG tag] [BATCHSIZE n [MINBATCHSIZE m]] [INPUTS name1 name2 ... OUTPUTS name1 name2 ...] model_blob
+* AI.MODELSET model_key backend device [TAG tag] [BATCHSIZE n [MINBATCHSIZE m]] [INPUTS name1 name2 ... OUTPUTS name1 name2 ...] BLOB model_blob
 */
 int RedisAI_ModelSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   RedisModule_AutoMemory(ctx);
@@ -121,7 +121,14 @@ int RedisAI_ModelSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
   const char* devicestr;
   AC_GetString(&ac, &devicestr, NULL, 0); 
 
-  if (strlen(devicestr) > 10) {
+  if (strlen(devicestr) > 10 ||
+      strcasecmp(devicestr, "INPUTS") == 0 ||
+      strcasecmp(devicestr, "OUTPUTS") == 0 ||
+      strcasecmp(devicestr, "TAG") == 0 ||
+      strcasecmp(devicestr, "BATCHSIZE") == 0 ||
+      strcasecmp(devicestr, "MINBATCHSIZE") == 0 ||
+      strcasecmp(devicestr, "BLOB") == 0
+      ) {
     return RedisModule_ReplyWithError(ctx, "ERR Invalid DEVICE");
   }
 
@@ -150,21 +157,21 @@ int RedisAI_ModelSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
     }
   }
 
-
   if (AC_IsAtEnd(&ac)) {
     return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, missing model BLOB");
   }
 
   ArgsCursor optionsac;
-  AC_GetSliceToOffset(&ac, &optionsac, argc-2);
+  const char* blob_matches[] = {"BLOB"};
+  AC_GetSliceUntilMatches(&ac, &optionsac, 1, blob_matches);
 
   if (optionsac.argc == 0 && backend == RAI_BACKEND_TENSORFLOW) {
     return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, INPUTS and OUTPUTS not specified");
   }
 
   ArgsCursor inac = {0};
   ArgsCursor outac = {0};
-  if (optionsac.argc > 0) {
+  if (optionsac.argc > 0 && backend == RAI_BACKEND_TENSORFLOW) {
     if (!AC_AdvanceIfMatch(&optionsac, "INPUTS")) {
       return RedisModule_ReplyWithError(ctx, "ERR INPUTS not specified");
     }
@@ -202,9 +209,39 @@ int RedisAI_ModelSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
 
   RAI_Model *model = NULL;
 
+  AC_AdvanceUntilMatches(&ac, 1, blob_matches);
+
+  if (AC_IsAtEnd(&ac)) {
+    return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, missing model BLOB");
+  }
+
+  AC_Advance(&ac);
+
+  ArgsCursor blobsac;
+  AC_GetSliceToEnd(&ac, &blobsac);
+
   size_t modellen;
-  const char *modeldef;
-  AC_GetString(&ac, &modeldef, &modellen, 0); 
+  char *modeldef;
+
+  if (blobsac.argc == 1) {
+    AC_GetString(&blobsac, (const char**)&modeldef, &modellen, 0); 
+  }
+  else {
+    const char *chunks[blobsac.argc];
+    size_t chunklens[blobsac.argc];
+    modellen = 0;
+    while (!AC_IsAtEnd(&blobsac)) {
+      AC_GetString(&blobsac, &chunks[blobsac.offset], &chunklens[blobsac.offset], 0); 
+      modellen += chunklens[blobsac.offset-1];
+    }
+    
+    modeldef = RedisModule_Calloc(modellen, sizeof(char));
+    size_t offset = 0;
+    for (size_t i=0; i<blobsac.argc; i++) {
+      memcpy(modeldef + offset, chunks[i], chunklens[i]);
+      offset += chunklens[i];
+    }
+  }
 
   RAI_Error err = {0};
 
@@ -223,6 +260,10 @@ int RedisAI_ModelSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
     model = RAI_ModelCreate(backend, devicestr, tag, opts, ninputs, inputs, noutputs, outputs, modeldef, modellen, &err);
   }
 
+  if (blobsac.argc > 1) {
+    RedisModule_Free(modeldef);
+  }
+
   if (err.code != RAI_OK) {
     #ifdef RAI_PRINT_BACKEND_ERRORS
     printf("ERR: %s\n", err.detail);
@@ -502,14 +543,14 @@ int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
   ArgsCursor outac = {0};
 
   if (!AC_AdvanceIfMatch(&ac, "INPUTS")) {
-    return RedisModule_ReplyWithError(ctx, "INPUTS not specified");
+    return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, INPUTS not specified");
   }
 
   const char* matches[] = {"OUTPUTS"};
   AC_GetSliceUntilMatches(&ac, &inac, 1, matches);
 
   if (!AC_AdvanceIfMatch(&ac, "OUTPUTS")) {
-    return RedisModule_ReplyWithError(ctx, "OUTPUTS not specified");
+    return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, OUTPUTS not specified");
   }
 
   AC_GetSliceToEnd(&ac, &outac);
@@ -541,15 +582,15 @@ int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
     if (!RAI_ScriptRunCtxAddInput(rinfo->sctx, t)) {
       RAI_FreeRunInfo(ctx,rinfo);
       RedisModule_CloseKey(key);
-      return RedisModule_ReplyWithError(ctx, "Input key not found");
+      return RedisModule_ReplyWithError(ctx, "ERR Input key not found");
     }
   }
 
   for (size_t i=0; i<noutputs; i++) {
     if (!RAI_ScriptRunCtxAddOutput(rinfo->sctx)) {
       RAI_FreeRunInfo(ctx,rinfo);
       RedisModule_CloseKey(key);
-      return RedisModule_ReplyWithError(ctx, "Output key not found");
+      return RedisModule_ReplyWithError(ctx, "ERR Output key not found");
     }
     RedisModule_RetainString(ctx, outputs[i]);
     array_append(rinfo->outkeys,outputs[i]);
@@ -651,12 +692,12 @@ int RedisAI_ScriptDel_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
 }
 
 /** 
-* AI.SCRIPTSET script_key device [TAG tag] script_source
+* AI.SCRIPTSET script_key device [TAG tag] SOURCE script_source
 */
 int RedisAI_ScriptSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   RedisModule_AutoMemory(ctx);
 
-  if (argc != 4 && argc != 6) return RedisModule_WrongArity(ctx);
+  if (argc != 5 && argc != 7) return RedisModule_WrongArity(ctx);
 
   ArgsCursor ac;
   ArgsCursor_InitRString(&ac, argv+1, argc-1);
@@ -673,14 +714,21 @@ int RedisAI_ScriptSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
   }
 
   if (AC_IsAtEnd(&ac)) {
-    return RedisModule_ReplyWithError(ctx, "Insufficient arguments, missing script definition");
+    return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, missing script SOURCE");
   }
 
-  RAI_Script *script = NULL;
-
   size_t scriptlen;
-  const char *scriptdef;
-  AC_GetString(&ac, &scriptdef, &scriptlen, 0); 
+  const char *scriptdef = NULL;
+
+  if (AC_AdvanceIfMatch(&ac, "SOURCE")) {
+    AC_GetString(&ac, &scriptdef, &scriptlen, 0); 
+  }
+
+  if (scriptdef == NULL) {
+    return RedisModule_ReplyWithError(ctx, "ERR Insufficient arguments, missing script SOURCE");
+  }
+
+  RAI_Script *script = NULL;
 
   RAI_Error err = {0};
   script = RAI_ScriptCreate(devicestr, tag, scriptdef, &err);

test/tests_dag.py

@@ -144,7 +144,7 @@ def test_dag_modelrun_financialNet_errors(env):
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
         env)
     ret = con.execute_command('AI.MODELSET', 'financialNet', 'TF', "CPU",
-                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', model_pb)
+                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
 
     tensor_number=1
@@ -390,7 +390,7 @@ def test_dag_modelrun_financialNet_separate_tensorget(env):
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
         env)
     ret = con.execute_command('AI.MODELSET', 'financialNet', 'TF', "CPU",
-                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', model_pb)
+                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
 
     tensor_number = 1
@@ -432,7 +432,7 @@ def test_dag_modelrun_financialNet(env):
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
         env)
     ret = con.execute_command('AI.MODELSET', 'financialNet', 'TF', "CPU",
-                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', model_pb)
+                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
 
     tensor_number = 1
@@ -471,7 +471,7 @@ def test_dag_modelrun_financialNet_no_writes(env):
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
         env)
     ret = con.execute_command('AI.MODELSET', 'financialNet', 'TF', "CPU",
-                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', model_pb)
+                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
 
     tensor_number = 1
@@ -522,7 +522,7 @@ def test_dagro_modelrun_financialNet_no_writes_multiple_modelruns(env):
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
         env)
     ret = con.execute_command('AI.MODELSET', 'financialNet', 'TF', DEVICE,
-                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', model_pb)
+                              'INPUTS', 'transaction', 'reference', 'OUTPUTS', 'output', 'BLOB', model_pb)
     env.assertEqual(ret, b'OK')
 
     tensor_number = 1