better error message

Author: hhsecond

redisai/command_builder.py

@@ -15,6 +15,10 @@ def modelset(self, name: AnyStr, backend: str, device: str, data: ByteString,
                  batch: int, minbatch: int, tag: AnyStr,
                  inputs: Union[AnyStr, List[AnyStr]],
                  outputs: Union[AnyStr, List[AnyStr]]) -> Sequence:
+        if device.upper() not in utils.allowed_devices:
+            raise ValueError(f"Device not allowed. Use any from {utils.allowed_devices}")
+        if backend.upper() not in utils.allowed_backends:
+            raise ValueError(f"Backend not allowed. Use any from {utils.allowed_backends}")
         args = ['AI.MODELSET', name, backend, device]
 
         if batch is not None:
@@ -87,6 +91,8 @@ def tensorget(self,
         return args
 
     def scriptset(self, name: AnyStr, device: str, script: str, tag: AnyStr = None) -> Sequence:
+        if device.upper() not in utils.allowed_devices:
+            raise ValueError(f"Device not allowed. Use any from {utils.allowed_devices}")
         args = ['AI.SCRIPTSET', name, device]
         if tag:
             args += ['TAG', tag]

redisai/utils.py

@@ -16,6 +16,9 @@
     'uint32': 'UINT32',
     'uint64': 'UINT64'}
 
+allowed_devices = {'CPU', 'GPU'}
+allowed_backends = {'TF', 'TFLITE', 'TORCH', 'ONNX'}
+
 
 def numpy2blob(tensor: np.ndarray) -> tuple:
     """Convert the numpy input from user to `Tensor`."""

test/test.py

@@ -106,6 +106,18 @@ def test_numpy_tensor(self):
         with self.assertRaises(TypeError):
             con.tensorset('trying', stringarr)
 
+    def test_modelset_errors(self):
+        model_path = os.path.join(MODEL_DIR, 'graph.pb')
+        model_pb = load_model(model_path)
+        con = self.get_client()
+        with self.assertRaises(ValueError):
+            con.modelset('m', 'tf', 'wrongdevice', model_pb,
+                         inputs=['a', 'b'], outputs=['mul'], tag='v1.0')
+        with self.assertRaises(ValueError):
+            con.modelset('m', 'wrongbackend', 'cpu', model_pb,
+                         inputs=['a', 'b'], outputs=['mul'], tag='v1.0')
+
+
     def test_modelget_meta(self):
         model_path = os.path.join(MODEL_DIR, 'graph.pb')
         model_pb = load_model(model_path)