Refactor conda backup script (#1350)

* Refacto conda backup

* Fix space

* Minor style

Author: atalman

s3_management/backup_conda.py

@@ -1,14 +1,16 @@
 #!/usr/bin/env python3
-# Downloads domain library packages from channel
+# Downloads domain pytorch and library packages from channel
 # And backs them up to S3
 # Do not use unless you know what you are doing
+# Usage:  python backup_conda.py --version 1.6.0
 
 import conda.api
 import boto3
 from typing import List, Optional
 import urllib
 import os
 import hashlib
+import argparse
 
 S3 = boto3.resource('s3')
 BUCKET = S3.Bucket('pytorch-backup')
@@ -23,11 +25,13 @@ def compute_md5(path:str) -> str:
 def download_conda_package(package:str, version:Optional[str] = None, depends:Optional[str] = None, channel:Optional[str] = None) -> List[str]:
     packages = conda.api.SubdirData.query_all(package, channels = [channel] if channel is not None else None, subdirs = _known_subdirs)
     rc = []
+
     for pkg in packages:
         if version is not None and pkg.version != version:
             continue
         if depends is not None and depends not in pkg.depends:
             continue
+
         print(f"Downloading {pkg.url}...")
         os.makedirs(pkg.subdir, exist_ok = True)
         fname = f"{pkg.subdir}/{pkg.fn}"
@@ -50,6 +54,18 @@ def upload_to_s3(prefix: str, fnames: List[str]) -> None:
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--version",
+        help="PyTorch Version to backup",
+        type=str,
+        required = True
+    )
+    options = parser.parse_args()
+    rc = download_conda_package("pytorch", channel = "pytorch", version = options.version)
+    upload_to_s3(f"v{options.version}/conda", rc)
+
     for libname in ["torchvision", "torchaudio", "torchtext"]:
-        rc = download_conda_package(libname, channel = "pytorch", depends = "pytorch 1.9.0")
-        upload_to_s3("v1.9.0-rc4/conda", rc)
+        print(f"processing {libname}")
+        rc = download_conda_package(libname, channel = "pytorch", depends = f"pytorch {options.version}")
+        upload_to_s3f(f"v{options.version}/conda", rc)