Add dagexecute support

pyproject.toml

@@ -41,6 +41,7 @@ tox-poetry = "^0.3.0"
 bandit = "^1.7.0"
 pylint = "^2.8.2"
 vulture = "^2.3"
+scikit-image = "==0.16.2"
 
 [tool.poetry.urls]
 url = "https://redisai.io"

redisai/client.py

@@ -70,7 +70,12 @@ def pipeline(self, transaction: bool = True, shard_hint: bool = None) -> "Pipeli
         )
 
     def dag(
-        self, load: Sequence = None, persist: Sequence = None, readonly: bool = False
+        self,
+        load: Sequence = None,
+        persist: Sequence = None,
+        routing: AnyStr = None,
+        timeout: int = None,
+        readonly: bool = False
     ) -> "Dag":
         """
         It returns a DAG object on which other DAG-allowed operations can be called. For
@@ -81,7 +86,16 @@ def dag(
         load : Union[AnyStr, List[AnyStr]]
             Load the list of given values from the keyspace to DAG scope
         persist : Union[AnyStr, List[AnyStr]]
-            Write the list of given key, values to the keyspace from DAG scope
+            For each tensor key in the given list, write its values to the keyspace from
+            DAG scope after the DAG execution is finished.
+        routing : AnyStr
+            Denotes a key to be used in the DAG 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.
+        timeout : int
+            The max number on milisecinds that may pass before the request is prossced
+            (meaning that the result will not be computed after that time and TIMEDOUT
+            is returned in that case)
         readonly : bool
             If True, it triggers AI.DAGRUN_RO, the read only DAG which cannot write (PERSIST) to
             the keyspace. But since it can't write, it can execute on replicas
@@ -105,9 +119,7 @@ def dag(
         >>> # You can even chain the operations
         >>> result = dag.tensorset(**akwargs).modelrun(**bkwargs).tensorget(**ckwargs).run()
         """
-        return Dag(
-            load, persist, self.execute_command, readonly, self.enable_postprocess
-        )
+        return Dag(load, persist, routing, timeout, self.execute_command, readonly)
 
     def loadbackend(self, identifier: AnyStr, path: AnyStr) -> str:
         """

redisai/dag.py

@@ -5,36 +5,57 @@
 
 from redisai import command_builder as builder
 from redisai.postprocessor import Processor
+from deprecated import deprecated
+import warnings
 
 processor = Processor()
 
 
 class Dag:
-    def __init__(self, load, persist, executor, readonly=False, postprocess=True):
+    def __init__(self, load, persist, routing, timeout, executor, readonly=False):
         self.result_processors = []
-        self.enable_postprocess = postprocess
-        if readonly:
-            if persist:
-                raise RuntimeError(
-                    "READONLY requests cannot write (duh!) and should not "
-                    "have PERSISTing values"
-                )
-            self.commands = ["AI.DAGRUN_RO"]
+        self.enable_postprocess = True
+        self.deprecatedDagrunMode = load is None and persist is None and routing is None
+        self.readonly = readonly
+        self.executor = executor
+
+        if readonly and persist:
+            raise RuntimeError(
+                "READONLY requests cannot write (duh!) and should not "
+                "have PERSISTing values"
+            )
+
+        if self.deprecatedDagrunMode:
+            # Throw warning about using deprecated dagrun
+            warnings.warn("Creating Dag without any of LOAD, PERSIST and ROUTING arguments"
+                          "is allowed only in deprecated AI.DAGRUN or AI.DAGRUN_RO commands", DeprecationWarning)
+            # Use dagrun
+            if readonly:
+                self.commands = ["AI.DAGRUN_RO"]
+            else:
+                self.commands = ["AI.DAGRUN"]
         else:
-            self.commands = ["AI.DAGRUN"]
-        if load:
+            # Use dagexecute
+            if readonly:
+                self.commands = ["AI.DAGEXECUTE_RO"]
+            else:
+                self.commands = ["AI.DAGEXECUTE"]
+        if load is not None:
             if not isinstance(load, (list, tuple)):
                 self.commands += ["LOAD", 1, load]
             else:
                 self.commands += ["LOAD", len(load), *load]
-        if persist:
+        if persist is not None:
             if not isinstance(persist, (list, tuple)):
-                self.commands += ["PERSIST", 1, persist, "|>"]
+                self.commands += ["PERSIST", 1, persist]
             else:
-                self.commands += ["PERSIST", len(persist), *persist, "|>"]
-        else:
-            self.commands.append("|>")
-        self.executor = executor
+                self.commands += ["PERSIST", len(persist), *persist]
+        if routing is not None:
+            self.commands += ["ROUTING", routing]
+        if timeout is not None:
+            self.commands += ["TIMEOUT", timeout]
+
+        self.commands.append("|>")
 
     def tensorset(
         self,
@@ -69,20 +90,71 @@ def tensorget(
         )
         return self
 
+    @deprecated(version="1.2.0", reason="Use modelexecute instead")
     def modelrun(
+            self,
+            key: AnyStr,
+            inputs: Union[AnyStr, List[AnyStr]],
+            outputs: Union[AnyStr, List[AnyStr]],
+    ) -> Any:
+        if self.deprecatedDagrunMode:
+            args = builder.modelrun(key, inputs, outputs)
+            self.commands.extend(args)
+            self.commands.append("|>")
+            self.result_processors.append(bytes.decode)
+            return self
+        else:
+            return self.modelexecute(key, inputs, outputs)
+
+    def modelexecute(
         self,
         key: AnyStr,
         inputs: Union[AnyStr, List[AnyStr]],
         outputs: Union[AnyStr, List[AnyStr]],
     ) -> Any:
-        args = builder.modelrun(key, inputs, outputs)
+        if self.deprecatedDagrunMode:
+            raise RuntimeError(
+                "You are using deprecated version of DAG, that does not supports MODELEXECUTE."
+                "The new version requires giving at least one of LOAD, PERSIST and ROUTING"
+                "arguments when constructing the Dag"
+            )
+        args = builder.modelexecute(key, inputs, outputs, None)
         self.commands.extend(args)
         self.commands.append("|>")
         self.result_processors.append(bytes.decode)
         return self
 
+    def scriptexecute(
+        self,
+        key: AnyStr,
+        function: str,
+        keys: Union[AnyStr, Sequence[AnyStr]] = None,
+        inputs: Union[AnyStr, Sequence[AnyStr]] = None,
+        args: Union[AnyStr, Sequence[AnyStr]] = None,
+        outputs: Union[AnyStr, List[AnyStr]] = None,
+    ) -> Any:
+        if self.readonly:
+            raise RuntimeError(
+                "AI.SCRIPTEXECUTE cannot be used in readonly mode"
+            )
+        if self.deprecatedDagrunMode:
+            raise RuntimeError(
+                "You are using deprecated version of DAG, that does not supports SCRIPTEXECUTE."
+                "The new version requires giving at least one of LOAD, PERSIST and ROUTING"
+                "arguments when constructing the Dag"
+            )
+        args = builder.scriptexecute(key, function, keys, inputs, args, outputs, None)
+        self.commands.extend(args)
+        self.commands.append("|>")
+        self.result_processors.append(bytes.decode)
+        return self
+
+    @deprecated(version="1.2.0", reason="Use execute instead")
     def run(self):
-        commands = self.commands[:-1]  # removing the last "|>
+        return self.execute()
+
+    def execute(self):
+        commands = self.commands[:-1]  # removing the last "|>"
         results = self.executor(*commands)
         if self.enable_postprocess:
             out = []