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()