added DAG ROUTING keyword

Author: DvirDukhan

docs/commands.md

@@ -820,7 +820,7 @@ It accepts one or more operations, split by the pipe-forward operator (`|>`).
 
 By default, the DAG execution context is local, meaning that tensor keys appearing in the DAG only live in the scope of the command. That is, setting a tensor with `TENSORSET` will store it local memory and not set it to an actual database key. One can refer to that key in subsequent commands within the DAG, but that key won't be visible outside the DAG or to other clients - no keys are open at the database level.
 
-Loading and persisting tensors from/to keyspace should be done explicitly. The user should specify which key tensors to load from keyspace using the `LOAD` keyword, and which command outputs to persist to the keyspace using the `PERSIST` keyspace. The user can also specify keys in Redis that are going to be accessed for read/write operations (for example, from within `AI.SCRIPTEXECUTE` command), by using the keyword `KEYS`.  
+Loading and persisting tensors from/to keyspace should be done explicitly. The user should specify which key tensors to load from keyspace using the `LOAD` keyword, and which command outputs to persist to the keyspace using the `PERSIST` keyspace. The user can also specify a tag or key which will assist for the routing of the DAG execution on the right shard in Redis that are going to be accessed for read/write operations (for example, from within `AI.SCRIPTEXECUTE` command), by using the keyword `ROUTING`.  
 
 As an example, if `command 1` sets a tensor, it can be referenced by any further command on the chaining.
 
@@ -832,7 +832,7 @@ A `TIMEOUT t` argument can be specified to cause a request to be removed from th
 ```
 AI.DAGEXECUTE [[LOAD <n> <key-1> <key-2> ... <key-n>] |
           [PERSIST <n> <key-1> <key-2> ... <key-n>] |
-          [KEYS <n> <key-1> <key-2> ... <key-n>]]+
+          [ROUTING <routing_tag>]]+
           [TIMEOUT t]
           |> <command> [|>  command ...]
 ```
@@ -841,9 +841,9 @@ _Arguments_
 
 * **LOAD**: denotes the beginning of the input tensors keys' list, followed by the number of keys, and one or more key names
 * **PERSIST**: denotes the beginning of the output tensors keys' list, followed by the number of keys, and one or more key names
-* **KEYS**: denotes the beginning of keys' list which are used within this command, followed by the number of keys, and one or more key names. Alternately, the keys names list can be replaced with a tag which all of those keys share. Redis will verify that all potential key accesses are done to the right shard.
+* **ROUTING**: denotes the a key name or a tag that will assist in routing the dag execution command to the right shard. Redis will verify that all potential key accesses are done to within the target shard.
 
-_While each of the LOAD, PERSIST and KEYS sections are optional (and may appear at most once in the command), the command must contain **at least one** of these 3 keywords._
+_While each of the LOAD, PERSIST and ROUTING sections are optional (and may appear at most once in the command), the command must contain **at least one** of these 3 keywords._
 * **TIMEOUT**: an optional argument, denotes the time (in ms) after which the client is unblocked and a `TIMEDOUT` string is returned
 * **|> command**: the chaining operator, that denotes the beginning of a RedisAI command, followed by one of RedisAI's commands. Command splitting is done by the presence of another `|>`. The supported commands are:
     * `AI.TENSORSET`
@@ -873,17 +873,17 @@ redis> AI.DAGEXECUTE PERSIST 1 predictions{tag} |>
 1) OK
 2) OK
 3) 1) FLOAT
-   2) 1) (integer) 2
-      2) (integer) 2
-   3) "\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x80@"
+   1) 1) (integer) 2
+      1) (integer) 2
+   2) "\x00\x00\x80?\x00\x00\x00@\x00\x00@@\x00\x00\x80@"
 ```
 
 A common pattern is enqueuing multiple SCRIPTEXECUTE and MODELEXECUTE commands within a DAG. The following example uses ResNet-50,to classify images into 1000 object categories. Given that our input tensor contains each color represented as a 8-bit integer and that neural networks usually work with floating-point tensors as their input we need to cast a tensor to floating-point and normalize the values of the pixels - for that we will use `pre_process_3ch` function. 
 
 To optimize the classification process we can use a post process script to return only the category position with the maximum classification - for that we will use `post_process` script. Using the DAG capabilities we've removed the necessity of storing the intermediate tensors in the keyspace. You can even run the entire process without storing the output tensor, as follows:
 
 ```
-redis> AI.DAGEXECUTE KEYS 1 {tag} |> 
+redis> AI.DAGEXECUTE ROUTING {tag} |> 
             AI.TENSORSET image UINT8 224 224 3 BLOB b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00....' |> 
             AI.SCRIPTEXECUTE imagenet_script{tag} pre_process_3ch INPUTS 1 image OUTPUTS 1 temp_key1 |> 
             AI.MODELEXECUTE imagenet_model{tag} INPUTS 1 temp_key1 OUTPUTS 1 temp_key2 |> 

src/execution/DAG/dag.c

@@ -667,8 +667,7 @@ int RedisAI_DagExecute_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
     size_t argpos = 1;
     while (argpos < argc) {
         const char *arg_string = RedisModule_StringPtrLen(argv[argpos++], NULL);
-        if (!strcasecmp(arg_string, "LOAD") || !strcasecmp(arg_string, "PERSIST") ||
-            !strcasecmp(arg_string, "KEYS")) {
+        if (!strcasecmp(arg_string, "LOAD") || !strcasecmp(arg_string, "PERSIST")) {
             if (argpos >= argc) {
                 return REDISMODULE_ERR;
             }
@@ -684,6 +683,11 @@ int RedisAI_DagExecute_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
             for (; argpos < last_argpos; argpos++) {
                 RedisModule_KeyAtPos(ctx, argpos);
             }
+        } else if (!strcasecmp(arg_string, "ROUTING")) {
+            if (argpos >= argc) {
+                return REDISMODULE_ERR;
+            }
+            RedisModule_KeyAtPos(ctx, argpos++);
         } else if (!strcasecmp(arg_string, "AI.MODELEXECUTE") ||
                    !strcasecmp(arg_string, "AI.SCRIPTEXECUTE")) {
             if (argpos >= argc) {

src/execution/parsing/dag_parser.c

@@ -245,7 +245,7 @@ int DAGInitialParsing(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleSt
     bool load_complete = false;
     bool persist_complete = false;
     bool timeout_complete = false;
-    bool keys_complete = false;
+    bool routing_complete = false;
 
     // The first arg is "AI.DAGEXECUTE(_RO) (or deprecated AI.DAGRUN(_RO))", so we go over from the
     // next arg.
@@ -280,13 +280,20 @@ int DAGInitialParsing(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleSt
             persist_complete = true;
             continue;
         }
-        if (!strcasecmp(arg_string, "KEYS") && !keys_complete && chainingOpCount == 0) {
-            const int parse_result =
-                ValidateKeysArgs(ctx, &argv[arg_pos], argc - arg_pos, rinfo->err);
-            if (parse_result <= 0)
+        if (!strcasecmp(arg_string, "ROUTING") && !routing_complete && chainingOpCount == 0) {
+            arg_pos++;
+            if (arg_pos == argc) {
+                RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR Missing ROUTING value");
                 return REDISMODULE_ERR;
-            arg_pos += parse_result;
-            keys_complete = true;
+            }
+            if (!VerifyKeyInThisShard(ctx, argv[arg_pos++])) {
+                RAI_SetError(
+                    rinfo->err, RAI_EDAGBUILDER,
+                    "ERR ROUTING value specified in the command hash to slot which does not "
+                    "belong to the current shard");
+                return REDISMODULE_ERR;
+            }
+            routing_complete = true;
             continue;
         }
         if (!strcasecmp(arg_string, "TIMEOUT") && !timeout_complete && chainingOpCount == 0) {
@@ -318,10 +325,10 @@ int DAGInitialParsing(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleSt
     // commands).
     if (!strncasecmp(RedisModule_StringPtrLen(argv[0], NULL), "AI.DAGEXECUTE",
                      strlen("AI.DAGEXECUTE"))) {
-        if (!load_complete && !persist_complete && !keys_complete) {
+        if (!load_complete && !persist_complete && !routing_complete) {
             RAI_SetError(rinfo->err, RAI_EDAGBUILDER,
                          "ERR AI.DAGEXECUTE and AI.DAGEXECUTE_RO commands must "
-                         "contain at least one out of KEYS, LOAD, PERSIST keywords");
+                         "contain at least one out of ROUTING, LOAD, PERSIST keywords");
             return REDISMODULE_ERR;
         }
     }
@@ -335,9 +342,9 @@ int DAGInitialParsing(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleSt
 int ParseDAGExecuteCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleString **argv,
                            int argc, bool dag_ro) {
 
-    // The minimal command is of the form: AI.DAGEXECUTE(_RO) KEYS/LOAD/PERSIST 1 <key> |>
+    // The minimal command is of the form: AI.DAGEXECUTE(_RO) ROUTING/LOAD/PERSIST 1 <key> |>
     // AI.TENSORGET <key>
-    if (argc < 7) {
+    if (argc < 6) {
         if (dag_ro) {
             RAI_SetError(rinfo->err, RAI_EDAGBUILDER,
                          "ERR missing arguments for 'AI.DAGEXECUTE_RO' command");

src/execution/parsing/script_commands_parser.c

@@ -147,17 +147,25 @@ static int _ScriptExecuteCommand_ParseCommand(RedisModuleCtx *ctx, RedisModuleSt
     if (keysRequired) {
         const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL);
         if (!strcasecmp(arg_string, "KEYS")) {
-            keysDone = true;
             argpos++;
+            keysDone = true;
             if (_ScriptExecuteCommand_ParseKeys(ctx, argv, argc, &argpos, error, sctx) ==
                 REDISMODULE_ERR) {
                 return REDISMODULE_ERR;
             }
+        } else if (!strcasecmp(arg_string, "INPUTS")) {
+            argpos++;
+            inputsDone = true;
+            if (_ScriptExecuteCommand_ParseInputs(ctx, argv, argc, &argpos, error, inputs) ==
+                REDISMODULE_ERR) {
+                return REDISMODULE_ERR;
+            }
         }
-        // argv[3] is not KEYS in AI.SCRIPTEXECUTE command (i.e., not in a DAG).
+        // argv[3] is not KEYS or INPUTS in AI.SCRIPTEXECUTE command (i.e., not in a DAG).
         else {
-            RAI_SetError(error, RAI_ESCRIPTRUN,
-                         "ERR KEYS scope must be provided first for AI.SCRIPTEXECUTE command");
+            RAI_SetError(
+                error, RAI_ESCRIPTRUN,
+                "ERR KEYS or INPUTS scope must be provided first for AI.SCRIPTEXECUTE command");
             return REDISMODULE_ERR;
         }
     }

tests/flow/tests_commands.py

@@ -161,13 +161,9 @@ def test_keys_syntax(env):
                                   "of given arguments",
                         "AI.SCRIPTEXECUTE script{1} bar KEYS 2 key{1}")
 
-    # ERR KEYS missing in AI.SCRIPTEXECUTE
-    check_error_message(env, con, "KEYS scope must be provided first for AI.SCRIPTEXECUTE command",
-                        "AI.SCRIPTEXECUTE script{1} bar INPUTS 2 a{1} a{1}")
-
-    # # ERR KEYS section in an inner AI.SCRIPTEXEUTE command within a DAG is not allowed.
-    # check_error_message(env, con, "Already encountered KEYS scope in current command",
-    #                     "AI.DAGEXECUTE KEYS 1 a{1} |> AI.SCRIPTEXECUTE script{1} bar KEYS 1 a{1}")
+    # ERR KEYS or INPUTS missing in AI.SCRIPTEXECUTE
+    check_error_message(env, con, "KEYS or INPUTS scope must be provided first for AI.SCRIPTEXECUTE command",
+                        "AI.SCRIPTEXECUTE script{1} bar OUTPUTS 2 a{1} a{1}")
 
 
 def test_scriptstore(env):
@@ -293,7 +289,7 @@ def test_pytorch_scriptexecute_errors(env):
 
     check_error_message(env, con, "Invalid arguments provided to AI.SCRIPTEXECUTE", 'AI.SCRIPTEXECUTE', 'ket{1}', 'bar', 'KEYS', 1, '{1}', 'ARGS')
 
-    check_error_message(env, con, "KEYS scope must be provided first for AI.SCRIPTEXECUTE command", 'AI.SCRIPTEXECUTE', 'ket{1}', 'bar', 'INPUTS', 'OUTPUTS')
+    check_error_message(env, con, "Invalid argument for inputs count in AI.SCRIPTEXECUTE", 'AI.SCRIPTEXECUTE', 'ket{1}', 'bar', 'INPUTS', 'OUTPUTS')
 
     check_error_message(env, con, "Invalid value for TIMEOUT",'AI.SCRIPTEXECUTE', 'ket{1}', 'bar', 'KEYS', 1, '{1}', 'INPUTS', 2, 'a{1}', 'b{1}', 'OUTPUTS', 1, 'c{1}', 'TIMEOUT', 'TIMEOUT')
 

tests/flow/tests_dag.py

@@ -388,15 +388,15 @@ def test_dagexecute_modelexecute_multidevice_resnet(env):
     ensureSlaveSynced(con, env)
 
     check_error_message(env, con, "INPUT key cannot be found in DAG",
-                        'AI.DAGEXECUTE', 'KEYS', '1', image_key, '|>', 'AI.SCRIPTEXECUTE',  script_name, 'pre_process_3ch',
+                        'AI.DAGEXECUTE', 'ROUTING', image_key, '|>', 'AI.SCRIPTEXECUTE',  script_name, 'pre_process_3ch',
                         'INPUTS', 1, image_key, 'OUTPUTS', 1, temp_key1)
 
     check_error_message(env, con, "INPUT key cannot be found in DAG",
-                        'AI.DAGEXECUTE', 'KEYS', '1', image_key, '|>', 'AI.MODELEXECUTE', model_name_0,
+                        'AI.DAGEXECUTE', 'ROUTING',  image_key, '|>', 'AI.MODELEXECUTE', model_name_0,
                         'INPUTS', 1, image_key, 'OUTPUTS', 1, temp_key1)
 
     ret = con.execute_command('AI.DAGEXECUTE', 
-                            'KEYS', 1, '{1}','|>',
+                            'ROUTING', '{1}','|>',
                             'AI.TENSORSET', image_key, 'UINT8', img.shape[1], img.shape[0], 3, 'BLOB', img.tobytes(),'|>',
                             'AI.SCRIPTEXECUTE',  script_name, 'wrong_fn',
                             'INPUTS', 1, image_key,
@@ -406,7 +406,7 @@ def test_dagexecute_modelexecute_multidevice_resnet(env):
     env.assertEquals("Function does not exist: wrong_fn",  ret[1].__str__())
 
     check_error_message(env, con, "Number of keys given as INPUTS here does not match model definition",
-                        'AI.DAGEXECUTE', 'KEYS', 1, '{1}',
+                        'AI.DAGEXECUTE', 'ROUTING', '{1}',
                         '|>', 'AI.TENSORSET', image_key, 'UINT8', img.shape[1], img.shape[0], 3, 'BLOB', img.tobytes(),
                         '|>',
                         'AI.SCRIPTEXECUTE',  script_name, 'pre_process_3ch',

tests/flow/tests_dag_basic.py

@@ -19,7 +19,7 @@ def test_dag_load(env):
 def test_dag_local_tensorset(env):
     con = env.getConnection()
 
-    command = "AI.DAGEXECUTE KEYS 1 {1} |> " \
+    command = "AI.DAGEXECUTE ROUTING {1} |> " \
               "AI.TENSORSET volatile_tensor1 FLOAT 1 2 VALUES 5 10 |> " \
               "AI.TENSORSET volatile_tensor2 FLOAT 1 2 VALUES 5 10 "
 
@@ -34,7 +34,7 @@ def test_dag_local_tensorset(env):
 def test_dagro_local_tensorset(env):
     con = env.getConnection()
 
-    command = "AI.DAGEXECUTE_RO KEYS 1 {some_tag} |> " \
+    command = "AI.DAGEXECUTE_RO ROUTING {some_tag} |> " \
               "AI.TENSORSET volatile_tensor1 FLOAT 1 2 VALUES 5 10 |> " \
               "AI.TENSORSET volatile_tensor2 FLOAT 1 2 VALUES 5 10 "
 
@@ -246,7 +246,7 @@ def test_dag_with_error(env):
     # Run the model from DAG context, where MODELEXECUTE op fails due to dim mismatch in one of the tensors inputs:
     # the input tensor 'b' is considered as tensor with dim 2X2X3 initialized with zeros, while the model expects that
     # both inputs to node 'mul' will be with dim 2.
-    ret = con.execute_command('AI.DAGEXECUTE_RO', 'KEYS', 1, '{1}',
+    ret = con.execute_command('AI.DAGEXECUTE_RO', 'ROUTING', '{1}',
                               '|>', 'AI.TENSORSET', 'a', 'FLOAT', 2, 'VALUES', 2, 3,
                               '|>', 'AI.TENSORSET', 'b', 'FLOAT', 2, 2, 3,
                               '|>', 'AI.MODELEXECUTE', 'tf_model{1}', 'INPUTS', 2, 'a', 'b', 'OUTPUTS', 1, 'tD',