Add package contraints to torchbench (#2318)

Summary:
This PR introduces package constraints because we want to guarantee that certain packages are the same version before and after the installation. This is important for the CI in pytorch/pytorch because overriding critical package versions will break the CI due to ABI incompatibility.

If we find a certain model breaks the constraints in the future, we could ask the model installer to apply the `constraint.txt` file with their requirements.txt file.

PyTorch CI Test: pytorch/pytorch#129114

Pull Request resolved: #2318

Reviewed By: huydhn

Differential Revision: D58801449

Pulled By: xuzhao9

fbshipit-source-id: 13e9c258f04fecd2b14e5a236c4c07ce54873d9b

install.py

@@ -5,26 +5,12 @@
 from pathlib import Path
 
 from userbenchmark import list_userbenchmarks
-from utils import get_pkg_versions, TORCH_DEPS
+from utils import get_pkg_versions, TORCH_DEPS, generate_pkg_constraints
+from utils.python_utils import pip_install_requirements
 
 REPO_ROOT = Path(__file__).parent
 
 
-def pip_install_requirements(requirements_txt="requirements.txt"):
-    try:
-        subprocess.run(
-            [sys.executable, "-m", "pip", "install", "-q", "-r", requirements_txt],
-            check=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-        )
-    except subprocess.CalledProcessError as e:
-        return (False, e.output)
-    except Exception as e:
-        return (False, e)
-    return True, None
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -54,6 +40,11 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
         action="store_true",
         help="Only require numpy to be installed, ignore torch, torchvision and torchaudio."
     )
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Only run the version check and generate the contraints"
+    )
     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")
@@ -71,10 +62,12 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
     if args.numpy:
         TORCH_DEPS = ["numpy"]
     print(
-        f"checking packages {', '.join(TORCH_DEPS)} are installed...",
+        f"checking packages {', '.join(TORCH_DEPS)} are installed, generating constaints...",
         end="",
         flush=True,
     )
+    if args.userbenchmark:
+        TORCH_DEPS = ["numpy", "torch"]
     try:
         versions = get_pkg_versions(TORCH_DEPS)
     except ModuleNotFoundError as e:
@@ -83,8 +76,12 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
             f"Error: Users must first manually install packages {TORCH_DEPS} before installing the benchmark."
         )
         sys.exit(-1)
+    generate_pkg_constraints(versions)
     print("OK")
 
+    if args.check_only:
+        exit(0)
+
     if args.userbenchmark:
         # Install userbenchmark dependencies if exists
         userbenchmark_dir = REPO_ROOT.joinpath("userbenchmark", args.userbenchmark)
@@ -96,7 +93,7 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
             )
         sys.exit(0)
 
-    success, errmsg = pip_install_requirements()
+    success, errmsg = pip_install_requirements(continue_on_fail=True)
     if not success:
         print("Failed to install torchbenchmark requirements:")
         print(errmsg)
@@ -120,7 +117,9 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
     new_versions = get_pkg_versions(TORCH_DEPS)
     if versions != new_versions:
         print(
-            f"The torch packages are re-installed after installing the benchmark deps. \
+            f"The numpy and torch package versions become inconsistent after installing the benchmark deps. \
                 Before: {versions}, after: {new_versions}"
         )
         sys.exit(-1)
+    else:
+        print(f"installed torchbench with package constraints: {versions}")

requirements.txt

@@ -15,10 +15,7 @@ transformers==4.38.1
 MonkeyType
 psutil
 pyyaml
-# We need to pin numpy version to the same as the torch testing environment
-# which still supports python 3.8
-numpy==1.21.2; python_version < '3.11'
-numpy==1.26.0; python_version >= '3.11'
+numpy
 opencv-python
 submitit
 pynvml

torchbenchmark/__init__.py

@@ -184,10 +184,10 @@ def setup(
             versions = get_pkg_versions(TORCH_DEPS)
         success, errmsg, stdout_stderr = _install_deps(model_path, verbose=verbose)
         if test_mode:
-            new_versions = get_pkg_versions(TORCH_DEPS, reload=True)
+            new_versions = get_pkg_versions(TORCH_DEPS)
             if versions != new_versions:
                 print(
-                    f"The torch packages are re-installed after installing the benchmark model {model_path}. \
+                    f"The numpy and torch packages are re-installed after installing the benchmark model {model_path}. \
                         Before: {versions}, after: {new_versions}"
                 )
                 sys.exit(-1)

torchbenchmark/canary_models/DALLE2_pytorch/install.py

@@ -2,6 +2,7 @@
 import patch
 import subprocess
 import sys
+from utils.python_utils import pip_install_requirements
 
 def patch_dalle2():
     import dalle2_pytorch
@@ -12,8 +13,8 @@ def patch_dalle2():
         print("Failed to patch dalle2_pytorch/dalle2_pytorch.py. Exit.")
         exit(1)
 
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+def pip_install_requirements_dalle2():
+    pip_install_requirements()
     # DALLE2_pytorch requires embedding-reader
     # https://github.com/lucidrains/DALLE2-pytorch/blob/00e07b7d61e21447d55e6d06d5c928cf8b67601d/setup.py#L34
     # embedding-reader requires an old version of pandas and pyarrow
@@ -22,5 +23,5 @@ def pip_install_requirements():
     subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-U', 'pandas', 'pyarrow'])
 
 if __name__ == '__main__':
-    pip_install_requirements()
+    pip_install_requirements_dalle2()
     patch_dalle2()

torchbenchmark/canary_models/codellama/install.py

@@ -1,6 +1,4 @@
 
-import subprocess
-import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
 

torchbenchmark/canary_models/fambench_dlrm/install.py

@@ -2,6 +2,7 @@
 import sys
 import subprocess
 from torchbenchmark import REPO_PATH
+from utils.python_utils import pip_install_requirements
 
 
 def update_fambench_submodule():
@@ -17,12 +18,6 @@ def update_fambench_submodule():
     subprocess.check_call(update_command, cwd=REPO_PATH)
 
 
-def pip_install_requirements():
-    subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"]
-    )
-
-
 if __name__ == "__main__":
     update_fambench_submodule()
     pip_install_requirements()

torchbenchmark/canary_models/fambench_xlmr/install.py

@@ -2,7 +2,7 @@
 import sys
 import subprocess
 from torchbenchmark import REPO_PATH
-
+from utils.python_utils import pip_install_requirements
 
 def update_fambench_submodule():
     "Update FAMBench submodule of the benchmark repo"
@@ -19,9 +19,7 @@ def update_fambench_submodule():
 
 def pip_install_requirements():
     try:
-        subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"]
-        )
+        pip_install_requirements()
         # pin fairseq version
         # ignore deps specified in requirements.txt
         subprocess.check_call(

torchbenchmark/canary_models/gat/install.py

@@ -1,13 +1,6 @@
-
-import subprocess
-import sys
 from utils import s3_utils
-
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt', '-f', 'https://data.pyg.org/whl/torch-2.1.0+cpu.html'])
-
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     s3_utils.checkout_s3_data("INPUT_TARBALLS", "Reddit_minimal.tar.gz", decompress=True)
-    pip_install_requirements()
+    pip_install_requirements(extra_args=["-f", "https://data.pyg.org/whl/torch-2.1.0+cpu.html"])

torchbenchmark/canary_models/hf_MPT_7b_instruct/install.py

@@ -1,10 +1,6 @@
-import subprocess
-import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/canary_models/hf_Yi/install.py

@@ -2,9 +2,7 @@
 import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/canary_models/hf_mixtral/install.py

@@ -1,10 +1,7 @@
-import subprocess
-import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
+from utils.python_utils import pip_install_requirements
 
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/canary_models/phi_1_5/install.py

@@ -2,9 +2,7 @@
 import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/canary_models/phi_2/install.py

@@ -1,10 +1,6 @@
-import subprocess
-import sys
 import os
 from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/canary_models/sage/install.py

@@ -1,12 +1,5 @@
-
-import subprocess
-import sys
 from utils import s3_utils
-
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt', '-f', 'https://data.pyg.org/whl/torch-2.1.0+cpu.html'])
-
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     s3_utils.checkout_s3_data("INPUT_TARBALLS", "Reddit_minimal.tar.gz", decompress=True)

torchbenchmark/canary_models/torchrec_dlrm/install.py

@@ -1,12 +1,4 @@
-import subprocess
-import sys
-
-
-def pip_install_requirements():
-    subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"]
-    )
-
+from utils.python_utils import pip_install_requirements
 
 if __name__ == "__main__":
     pip_install_requirements()

torchbenchmark/models/Background_Matting/install.py

@@ -1,10 +1,7 @@
-import subprocess
-import sys
-from utils import s3_utils
+from utils import s3_utils, python_utils
 
 def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip',
-                           'install', '-q', '-r', 'requirements.txt'])
+    python_utils.pip_install_requirements('requirements.txt')
 
 if __name__ == '__main__':
     pip_install_requirements()

torchbenchmark/models/Background_Matting/requirements.txt

@@ -1,7 +1,4 @@
-# We need to pin numpy version to the same as the torch testing environment
-# which still supports python 3.8
-numpy==1.21.2; python_version < '3.11'
-numpy==1.26.0; python_version >= '3.11'
+numpy
 opencv-python
 pandas
 Pillow

torchbenchmark/models/LearningToPaint/install.py

@@ -1,11 +1,7 @@
-import subprocess
-import sys
 from utils import s3_utils
+from utils.python_utils import pip_install_requirements
 
 
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
-
 if __name__ == '__main__':
     s3_utils.checkout_s3_data("INPUT_TARBALLS", "Super_SloMo_inputs.tar.gz", decompress=True)
     pip_install_requirements()

torchbenchmark/models/Super_SloMo/install.py

@@ -1,11 +1,5 @@
-import subprocess
-import sys
 from utils import s3_utils
-
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
-
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     s3_utils.checkout_s3_data("INPUT_TARBALLS", "Super_SloMo_inputs.tar.gz", decompress=True)

torchbenchmark/models/dcgan/install.py

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

torchbenchmark/models/demucs/install.py

@@ -1,17 +1,4 @@
-import subprocess
-import sys
-
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
-
-def spacy_download(language):
-    pass
-
-def preprocess():
-    pass
+from utils.python_utils import pip_install_requirements
 
 if __name__ == '__main__':
     pip_install_requirements()
-    spacy_download('')
-    preprocess()

torchbenchmark/models/dlrm/install.py

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