Merge pull request #1 from RedisAI/client-api

Client api

Author: mnunberg

redisai/__init__.py

@@ -1,3 +1,3 @@
-from .client import Client, Type
-
-
+from .client import (
+    Client, Tensor, ScalarTensor, BlobTensor, DType, Device, Backend
+)

redisai/_util.py

@@ -1,5 +1,6 @@
 import six
 
+
 def to_string(s):
     if isinstance(s, six.string_types):
         return s

redisai/client.py

@@ -1,40 +1,231 @@
-from redis import StrictRedis
-from redis._compat import (long, nativestr)
 from enum import Enum
-import six
-
-class Type(Enum):
-    FLOAT=1
-    DOUBLE=2
-    INT8=3
-    INT16=4
-    INT32=5
-    INT64=6
-    UINT8=7
-    UINT16=8
-    
+from redis import StrictRedis
+from ._util import to_string
+
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
+try:
+    from typing import Union, Any, AnyStr, ByteString, Collection
+except ImportError:
+    pass
+
+
+class Device(Enum):
+    cpu = 'cpu'
+    gpu = 'gpu'
+
+
+class Backend(Enum):
+    tf = 'tf'
+    torch = 'torch'
+    onnx = 'ort'
+
+
+class DType(Enum):
+    float = 'float'
+    double = 'double'
+    int8 = 'int8'
+    int16 = 'int16'
+    int32 = 'int32'
+    int64 = 'int64'
+    uint8 = 'uint8'
+    uint16 = 'uint16'
+    uint32 = 'uint32'
+    uint64 = 'uint64'
+
+    # aliases
+    float32 = 'float'
+    float64 = 'double'
 
-class Client(StrictRedis):
 
-    def __init__(self, *args, **kwargs):
+class Tensor(object):
+    ARGNAME = 'VALUES'
+
+    def __init__(self,
+                 dtype,  # type: DType
+                 shape,  # type: Collection[int]
+                 value):
+        """
+        Declare a tensor suitable for passing to tensorset
+        :param dtype: The type the values should be stored as.
+            This can be one of Tensor.FLOAT, tensor.DOUBLE, etc.
+        :param shape: An array describing the shape of the tensor. For an
+            image 250x250 with three channels, this would be [250, 250, 3]
+        :param value: The value for the tensor. Can be an array.
+            The contents must coordinate with the shape, meaning that the
+            overall length needs to be the product of all figures in the
+            shape. There is no verification to ensure that each dimension
+            is correct. Your application must ensure that the ordering
+            is always consistent.
         """
-        Create a new Client optional host and port
+        self.type = dtype
+        self.shape = shape
+        self.value = value
+        self._size = 1
+        if not isinstance(value, (list, tuple)):
+            self.value = [value]
+
+    @property
+    def size(self):
+        return self._size
+
+    def __repr__(self):
+        return '<{c.__class__.__name__}(shape={s} type={t}) at 0x{id:x}>'.format(
+            c=self,
+            s=self.shape,
+            t=self.type,
+            id=id(self))
+
+
+class ScalarTensor(Tensor):
+    def __init__(self, dtype, *values):
+        # type: (ScalarTensor, DType, Any) -> None
+        """
+        Declare a tensor with a bunch of scalar values. This can be used
+        to 'batch-load' several tensors.
+
+        :param dtype: The datatype to store the tensor as
+        :param values: List of values
+        """
+        super(ScalarTensor, self).__init__(dtype, [1], values)
+        self._size = len(values)
+
+
+class BlobTensor(Tensor):
+    ARGNAME = 'BLOB'
 
-        If conn is not None, we employ an already existing redis connection
+    def __init__(self,
+                 dtype,
+                 shape,  # type: Collection[int]
+                 *blobs  # type: Union[BlobTensor, ByteString]
+                 ):
         """
-        StrictRedis.__init__(self, *args, **kwargs)
-        
-                # Set the module commands' callbacks
-        MODULE_CALLBACKS = {
-                'AI.TENSORSET': lambda r: r and nativestr(r) == 'OK',
+        Create a tensor from a binary blob
+        :param dtype: The datatype, one of Tensor.FLOAT, Tensor.DOUBLE, etc.
+        :param shape: An array
+        :param blobs: One or more blobs to assign to the tensor.
+        """
+        if len(blobs) > 1:
+            blobarr = bytearray()
+            for b in blobs:
+                if isinstance(b, BlobTensor):
+                    b = b.value[0]
+                blobarr += b
+            size = len(blobs)
+            blobs = bytes(blobarr)
+        else:
+            blobs = bytes(blobs[0])
+            size = 1
+
+        super(BlobTensor, self).__init__(dtype, shape, blobs)
+        self._size = size
+
+    @classmethod
+    def from_numpy(cls, *nparrs):
+        # type: (type, np.array) -> BlobTensor
+        blobs = []
+        for arr in nparrs:
+            blobs.append(arr.data)
+        dt = DType.__members__[str(nparrs[0].dtype)]
+        return cls(dt, nparrs[0].shape, *blobs)
+
+    @property
+    def blob(self):
+        return self.value[0]
+
+    def to_numpy(self):
+        # type: () -> np.array
+        a = np.frombuffer(self.value[0], dtype=self._to_numpy_type(self.type))
+        return a.reshape(self.shape)
+
+    @staticmethod
+    def _to_numpy_type(t):
+        t = t.lower()
+        mm = {
+            'float': 'float32',
+            'double': 'float64'
+        }
+        if t in mm:
+            return mm[t]
+        return t
+
+
+class Client(StrictRedis):
+    def modelset(self,
+                 name,  # type: AnyStr
+                 backend,  # type: Backend
+                 device,  # type: Device
+                 inputs,  # type: Collection[AnyStr]
+                 outputs,  # type: Collection[AnyStr]
+                 data  # type: ByteString
+                 ):
+        args = ['AI.MODELSET', name, backend.value, device.value, 'INPUTS']
+        args += inputs
+        args += ['OUTPUTS'] + outputs
+        args += [data]
+        return self.execute_command(*args)
+
+    def modelget(self, name):
+        rv = self.execute_command('AI.MODELGET', name)
+        return {
+            'backend': Backend(rv[0]),
+            'device': Device(rv[1]),
+            'data': rv[2]
         }
-        for k, v in six.iteritems(MODULE_CALLBACKS):
-            self.set_response_callback(k, v)
 
+    def modelrun(self, name, inputs, outputs):
+        args = ['AI.MODELRUN', name]
+        args += ['INPUTS'] + inputs + ['OUTPUTS'] + outputs
+        return self.execute_command(*args)
 
-    def tensorset(self, key, type, dimensions, tensor):
-        args = ['AI.TENSORSET', key, type.name] + dimensions + ['VALUES'] + tensor
-        
+    def tensorset(self, key, tensor):
+        # type: (Client, AnyStr, Union[Tensor, np.ndarray]) -> Any
+        """
+        Set the values of the tensor on the server using the provided Tensor object
+        :param key: The name of the tensor
+        :param tensor: a `Tensor` object
+        """
+        if np and isinstance(tensor, np.ndarray):
+            tensor = BlobTensor.from_numpy(tensor)
+        args = ['AI.TENSORSET', key, tensor.type.value, tensor.size]
+        args += tensor.shape
+        args += [tensor.ARGNAME]
+        args += tensor.value
         return self.execute_command(*args)
 
- 
+    def tensorget(self, key, astype=Tensor, meta_only=False):
+        """
+        Retrieve the value of a tensor from the server
+        :param key: the name of the tensor
+        :param astype: the resultant tensor type
+        :param meta_only: if true, then the value is not retrieved,
+            only the shape and the type
+        :return: an instance of astype
+        """
+        argname = 'META' if meta_only else astype.ARGNAME
+        res = self.execute_command('AI.TENSORGET', key, argname)
+        dtype, shape = to_string(res[0]), res[1]
+        if meta_only:
+            return astype(dtype, shape, [])
+        else:
+            return astype(dtype, shape, res[2])
+
+    def scriptset(self, name, device, script):
+        return self.execute_command('AI.SCRIPTSET', name, device.value, script)
+
+    def scriptget(self, name):
+        r = self.execute_command('AI.SCRIPTGET', name)
+        return {
+            'device': to_string(r[0]),
+            'script': to_string(r[1])
+        }
+
+    def scriptrun(self, name, function, inputs, outputs):
+        args = ['AI.SCRIPTRUN', name, function, 'INPUTS']
+        args += inputs
+        args += ['OUTPUTS']
+        args += outputs
+        return self.execute_command(*args)