Fixes to device allocation to make CPU tests pass

Author: lantiga

src/backends/tensorflow.c

@@ -515,6 +515,18 @@ int RAI_ModelRunTF(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
         }
     }
 
+    char tf_devicestr[256];
+    int devicestr_len = strlen(mctxs[0]->model->devicestr);
+    if (strncasecmp(mctxs[0]->model->devicestr, "CPU", 3) == 0) {
+        sprintf(tf_devicestr, "/device:CPU:0");
+    }
+    else if (devicestr_len == 3) {
+        sprintf(tf_devicestr, "/device:%s:0", mctxs[0]->model->devicestr);
+    }
+    else {
+        sprintf(tf_devicestr, "/device:%s", mctxs[0]->model->devicestr);
+    }
+
     for (size_t i = 0; i < ninputs; ++i) {
         RAI_Tensor *batched_input_tensors[nbatches];
 
@@ -530,11 +542,19 @@ int RAI_ModelRunTF(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
             RedisModule_Free(errorMessage);
             return 1;
         }
-        // TODO EAGER
+
         inputTensorsHandles[i] = TFE_TensorHandleCopyToDevice(inputTensorsHandles[i],
                                                               mctxs[0]->model->session,
-                                                              mctxs[0]->model->devicestr,
+                                                              tf_devicestr,
                                                               status);
+
+        if (TF_GetCode(status) != TF_OK) {
+            char *errorMessage = RedisModule_Strdup(TF_Message(status));
+            RAI_SetError(error, RAI_EMODELRUN, errorMessage);
+            TF_DeleteStatus(status);
+            RedisModule_Free(errorMessage);
+            return 1;
+        }
     }
 
     TFE_Op *fn_op = TFE_NewOp(mctxs[0]->model->session, RAI_TF_FN_NAME, status);
@@ -555,8 +575,6 @@ int RAI_ModelRunTF(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
         return 1;
     }
 
-    // TODO EAGER: send tensors to device (as long as we keep device allocation EXPLICIT)
-
     int noutputs_ = noutputs;
     TFE_Execute(fn_op, outputTensorsHandles, &noutputs_, status);
     if (TF_GetCode(status) != TF_OK) {
@@ -583,7 +601,7 @@ int RAI_ModelRunTF(RAI_ModelRunCtx **mctxs, RAI_Error *error) {
     for (size_t i = 0; i < noutputs; ++i) {
         outputTensorsHandles[i] = TFE_TensorHandleCopyToDevice(outputTensorsHandles[i],
                                                                mctxs[0]->model->session,
-                                                               "CPU",
+                                                               "/device:CPU:0",
                                                                status);
 
         outputTensorsValues[i] = TFE_TensorHandleResolve(outputTensorsHandles[i], status);

tests/flow/test_serializations.py

@@ -81,7 +81,7 @@ def test_v0_torch_model(self):
         self.env.assertEqual([backend, device, tag, batchsize, minbatchsize, inputs, outputs], [b"TORCH", b"CPU", b"PT_MINIMAL_V0", 0, 0, [b'a', b'b'], [b'']])  
         torch_model_run(self.env, key_name)
 
-    def test_v0_troch_script(self):
+    def test_v0_torch_script(self):
         key_name = "torch_script{1}"
         con = self.env.getConnection()
         script_rdb = b'\x07\x81\x00\x8f\xd2\t\x12\x0fL\x00\x05\x04CPU\x00\x05\x10TORCH_SCRIPT_V0\x00\x05\xc3@W@i\x0fdef bar(a, b):\n  \x00\x0ereturn a + b\n\nd\x80 \x08_variadic@)\x12args : List[Tensor]\xe0\x06;  \x02[0] A`\t\x031]\n\x00\x00\t\x00\x0b\xee\x04\xe7\x11\xaez\x91'