Allow using device id (#91)

* Add the option to use device id (i.e., "cpu\gpu:<number>", instead of just "cpu/gpu")

Author: alonre24

redisai/command_builder.py

@@ -25,7 +25,10 @@ def modelstore(
 ) -> Sequence:
     if name is None:
         raise ValueError("Model name was not given")
-    if device.upper() not in utils.allowed_devices:
+
+    # device format should be: "CPU | GPU [:<num>]"
+    device_type = device.split(":")[0]
+    if device_type.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}")

test/test.py

@@ -342,7 +342,7 @@ def test_modelexecute_non_list_input_output(self):
         ret = con.modelexecute("m", ["a", "b"], "out")
         self.assertEqual(ret, "OK")
 
-    def test_nonasciichar(self):
+    def test_non_ascii_char(self):
         nonascii = "ĉ"
         model_path = os.path.join(MODEL_DIR, tf_graph)
         model_pb = load_model(model_path)
@@ -363,6 +363,21 @@ def test_nonasciichar(self):
         tensor = con.tensorget("c" + nonascii)
         self.assertTrue((np.allclose(tensor, [4.0, 9.0])))
 
+    def test_device_with_id(self):
+        model_path = os.path.join(MODEL_DIR, tf_graph)
+        model_pb = load_model(model_path)
+        con = self.get_client()
+        ret = con.modelstore(
+            "m",
+            "tf",
+            "cpu:1",
+            model_pb,
+            inputs=["a", "b"],
+            outputs=["mul"],
+            tag="v1.0",
+        )
+        self.assertEqual('OK', ret)
+
     def test_run_tf_model(self):
         model_path = os.path.join(MODEL_DIR, tf_graph)
         bad_model_path = os.path.join(MODEL_DIR, torch_graph)