Fix paths

Author: dk19y

test/integration/__init__.py

@@ -23,36 +23,41 @@
 cpu_sub_dir = 'model_cpu'
 gpu_sub_dir = 'model_gpu'
 eia_sub_dir = 'model_eia'
+code_sub_dir = 'code'
 
 model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
-mnist_cpu_script = os.path.join(model_cpu_dir, 'code/mnist.py')
+mnist_cpu_script = os.path.join(model_cpu_dir, code_sub_dir, 'mnist.py')
 model_cpu_tar = file_utils.make_tarfile(mnist_cpu_script,
                                         os.path.join(model_cpu_dir, "model.pth"),
-                                        model_cpu_dir)
+                                        model_cpu_dir,
+                                        True)
 
 model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
-mnist_1d_script = os.path.join(model_cpu_1d_dir, 'code/mnist_1d.py')
+mnist_1d_script = os.path.join(model_cpu_1d_dir, code_sub_dir, 'mnist_1d.py')
 model_cpu_1d_tar = file_utils.make_tarfile(mnist_1d_script,
                                            os.path.join(model_cpu_1d_dir, "model.pth"),
-                                           model_cpu_1d_dir)
+                                           model_cpu_1d_dir,
+                                           True)
 
 model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
-mnist_gpu_script = os.path.join(model_gpu_dir, 'code/mnist.py')
+mnist_gpu_script = os.path.join(model_gpu_dir, code_sub_dir, 'mnist.py')
 model_gpu_tar = file_utils.make_tarfile(mnist_gpu_script,
                                         os.path.join(model_gpu_dir, "model.pth"),
-                                        model_gpu_dir)
+                                        model_gpu_dir,
+                                        True)
 
 model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
 mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
 model_eia_tar = file_utils.make_tarfile(mnist_eia_script,
                                         os.path.join(model_eia_dir, "model.pth"),
                                         model_eia_dir)
 
-call_model_fn_once_script = os.path.join(model_cpu_dir, 'code/call_model_fn_once.py')
+call_model_fn_once_script = os.path.join(model_cpu_dir, code_sub_dir, 'call_model_fn_once.py')
 call_model_fn_once_tar = file_utils.make_tarfile(call_model_fn_once_script,
                                                  os.path.join(model_cpu_dir, "model.pth"),
                                                  model_cpu_dir,
-                                                 "model_call_model_fn_once.tar.gz")
+                                                 "model_call_model_fn_once.tar.gz",
+                                                 True)
 
 ROLE = 'dummy/unused-role'
 DEFAULT_TIMEOUT = 20

test/utils/file_utils.py

@@ -16,9 +16,12 @@
 import tarfile
 
 
-def make_tarfile(script, model, output_path, filename="model.tar.gz"):
+def make_tarfile(script, model, output_path, filename="model.tar.gz", preserve_script_path=False):
     output_filename = os.path.join(output_path, filename)
     with tarfile.open(output_filename, "w:gz") as tar:
-        tar.add(script, arcname=os.path.basename(script))
+        if(preserve_script_path):
+            tar.add(script, arcname=script)
+        else:
+            tar.add(script, arcname=os.path.basename(script))
         tar.add(model, arcname=os.path.basename(model))
     return output_filename