refactor: utils/prepare/upload.py (#1607)

Co-authored-by: Fedor Ignatov <[email protected]>

Author: vaskonov

utils/prepare/upload.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 
 import argparse
-import os
-import shutil
+import pathlib
 import tarfile
 from pathlib import Path
 
@@ -23,45 +22,48 @@
 from hashes import main
 
 
-def upload(config_in_file):
+def upload(config_in_file: str, tar_name: str, tar_output_dir: Path):
+    if not tar_output_dir.exists():
+        raise RuntimeError(f'A folder {tar_output_dir} does not exist')
+
+    print(f'Config: {config_in_file}')
+    if not Path(config_in_file).exists():
+        raise RuntimeError(f'A config {config_in_file} does not exist')
 
-    print(config_in_file)
     config_in = parse_config(config_in_file)
     config_in_file = find_config(config_in_file)
 
     model_path = Path(config_in['metadata']['variables']['MODEL_PATH']).expanduser()
-    models_path = Path(config_in['metadata']['variables']['MODELS_PATH']).expanduser()
     model_name, class_name = config_in_file.stem, config_in_file.parent.name
-    
-    if str(model_name) not in str(model_path):
-        raise(f'{model_name} is not the path of the {model_path}')
-    
-    arcname = str(model_path).split("models/")[1]
-    tar_path = models_path/model_name
-    tmp_folder = f'/tmp/'
-    tmp_tar = tmp_folder + f'{model_name}.tar.gz'
 
-    print("model_path", model_path)
-    print("class_name", class_name)
-    print("model_name", model_name)
-    
-    print("Start tarring")
-    archive = tarfile.open(tmp_tar, "w|gz")
-    archive.add(model_path, arcname=arcname)
-    archive.close()
+    if tar_name is None:
+        tar_name = f'{model_name}'
+        print(f'tar_name set to {tar_name}')
+
+    full_tar_name = tar_output_dir / f'{tar_name}.tar.gz'
+    if Path(full_tar_name).exists():
+        raise RuntimeError(f'An archive {Path(full_tar_name)} already exists')
+
+    print(f'model_path: {model_path}')
+    print(f'class_name: {class_name}')
+    print(f'model_name: {model_name}')
+    print(f'Start tarring to {full_tar_name}')
+    with tarfile.open(str(full_tar_name), "w|gz") as archive:
+        archive.add(model_path, arcname=pathlib.os.sep)
+
     print("Stop tarring")
+    print(f'Tar archive: {Path(full_tar_name)} has been created')
 
     print("Calculating hash")
-    main(tmp_tar)
-
-    print("tmp_tar", tmp_tar)
-    command = f'scp -r {tmp_folder}{model_name}* share.ipavlov.mipt.ru:/home/export/v1/{class_name}'
-    donwload_url = f'http://files.deeppavlov.ai/v1/{class_name}/{model_name}.tar.gz'
-    print(command, donwload_url, sep='\n')
+    main(full_tar_name)
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("config_in", help="path to a config", type=str)
+    parser.add_argument('-c', '--config_in', help='path to a config', type=str)
+    parser.add_argument('-n', '--tar_name', help='name of the tar archive (without tar.gz extension)',
+                        default=None, required=False, type=str)
+    parser.add_argument('-o', '--tar_output_dir', help='dir to save a tar archive', default='./',
+                        required=False, type=Path)
     args = parser.parse_args()
-    upload(args.config_in)
+    upload(args.config_in, args.tar_name, args.tar_output_dir)