Move distributed install to userbenchmark (#1824)

Summary:
Add install script to userbenchmark directory. Each userbenchmark can install their own dependencies.

Pull Request resolved: #1824

Reviewed By: FindHao

Differential Revision: D48268234

Pulled By: xuzhao9

fbshipit-source-id: d566bfa782e1a31fed456d5bcb38f3074e0284f5

install.py

@@ -3,7 +3,9 @@
 import os
 import sys
 from utils import TORCH_DEPS, proxy_suggestion, get_pkg_versions, _test_https
+from userbenchmark import list_userbenchmarks
 from pathlib import Path
+
 REPO_ROOT = Path(__file__).parent
 
 def pip_install_requirements(requirements_txt="requirements.txt"):
@@ -27,7 +29,7 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
     parser.add_argument("--canary", action="store_true", help="Install canary model.")
     parser.add_argument("--continue_on_fail", action="store_true")
     parser.add_argument("--verbose", "-v", action="store_true")
-    parser.add_argument("--component", choices=["distributed"], help="Install requirements for optional components.")
+    parser.add_argument("--userbenchmark", choices=list_userbenchmarks(), help="Install requirements for optional components.")
     args = parser.parse_args()
 
     os.chdir(os.path.realpath(os.path.dirname(__file__)))
@@ -41,13 +43,11 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
         sys.exit(-1)
     print("OK")
 
-    if args.component == "distributed":
-        success, errmsg = pip_install_requirements(requirements_txt="torchbenchmark/util/distributed/requirements.txt")
-        if not success:
-            print("Failed to install torchbenchmark distributed requirements:")
-            print(errmsg)
-            if not args.continue_on_fail:
-                sys.exit(-1)
+    if args.userbenchmark:
+        # Install userbenchmark dependencies if exists
+        userbenchmark_dir = REPO_ROOT.joinpath("userbenchmark", args.userbenchmark)
+        if userbenchmark_dir.joinpath("install.py").is_file():
+            subprocess.check_call([sys.executable, "install.py"], cwd=userbenchmark_dir.absolute())
         sys.exit(0)
 
     success, errmsg = pip_install_requirements()

userbenchmark/__init__.py

@@ -1,14 +1,15 @@
-import yaml
 from pathlib import Path
+from typing import List
 
 CURRENT_DIR = Path(__file__).parent
 
-def list_userbenchmarks():
+def list_userbenchmarks() -> List[str]:
     ub_dirs = [x for x in CURRENT_DIR.iterdir() if x.is_dir() and x.joinpath('__init__.py').exists() ]
     ub_names = list(map(lambda x: x.name, ub_dirs))
     return ub_names
 
 def get_ci_from_ub(ub_name):
+    import yaml
     ci_file = CURRENT_DIR.joinpath(ub_name).joinpath("ci.yaml")
     if not ci_file.exists():
         return None

userbenchmark/distributed/install.py

@@ -0,0 +1,8 @@
+import sys
+import subprocess
+
+def pip_install_requirements():
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+
+if __name__ == '__main__':
+    pip_install_requirements()