Append extra args to pip_install_requirements

Author: xuzhao9

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/maml_omniglot/install.py

@@ -1,10 +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__':
     pip_install_requirements()

torchbenchmark/models/moondream/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/models/nvidia_deeprecommender/install.py

@@ -1,9 +1,6 @@
 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/opacus_cifar10/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()

torchbenchmark/models/pytorch_CycleGAN_and_pix2pix/install.py

@@ -1,10 +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", "pytorch_CycleGAN_and_pix2pix_inputs.tar.gz", decompress=True)

torchbenchmark/models/pytorch_stargan/install.py

@@ -1,9 +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", "pytorch_stargan_inputs.tar.gz", decompress=True)

torchbenchmark/models/pytorch_unet/install.py

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

torchbenchmark/models/sam/install.py

@@ -2,6 +2,7 @@
 import subprocess
 import sys
 import requests
+from utils.python_utils import pip_install_requirements
 
 def download(uri):
     directory = '.data'
@@ -16,10 +17,6 @@ def download(uri):
     else:
         print(f'Failed to download file with status code {response.status_code}')
 
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
-
 def download_checkpoint():
     download('https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth')
 

torchbenchmark/models/sam_fast/install.py

@@ -1,7 +1,6 @@
 import os
-import subprocess
-import sys
 import requests
+from utils.python_utils import pip_install_requirements
 
 def download(uri):
     directory = '.data'
@@ -16,10 +15,6 @@ def download(uri):
     else:
         print(f'Failed to download file with status code {response.status_code}')
 
-
-def pip_install_requirements():
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
-
 def download_checkpoint():
     download('https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth')
 

torchbenchmark/models/soft_actor_critic/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/speech_transformer/install.py

@@ -1,9 +1,5 @@
-import sys
-import subprocess
 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", "speech_transformer_inputs.tar.gz", decompress=True)

torchbenchmark/util/framework/lit_llama.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from torchbenchmark import REPO_PATH
+from utils.python_utils import pip_install_requirements
 
 LIT_LLAMA_PATH = os.path.join(REPO_PATH, "submodules", "lit-llama")
 
@@ -21,20 +22,6 @@ def update_lit_llama_submodule():
     subprocess.check_call(update_command, cwd=REPO_PATH)
 
 
-def pip_install_requirements():
-    subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "-q",
-            "-r",
-            os.path.join(LIT_LLAMA_PATH, "requirements.txt"),
-        ]
-    )
-
-
 def openllama_download():
     if os.path.exists(
         os.path.join(LIT_LLAMA_PATH, "checkpoints/lit-llama/7B/lit-llama.pth")

utils/python_utils.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 import subprocess
 
+from typing import Optional, List
+
 DEFAULT_PYTHON_VERSION = "3.11"
 
 PYTHON_VERSION_MAP = {
@@ -23,7 +25,10 @@ def create_conda_env(pyver: str, name: str):
     subprocess.check_call(command)
 
 
-def pip_install_requirements(requirements_txt="requirements.txt", continue_on_fail=False, no_build_isolation=False):
+def pip_install_requirements(requirements_txt="requirements.txt",
+                            continue_on_fail=False,
+                            no_build_isolation=False,
+                            extra_args: Optional[List[str]]=None):
     import sys
     constraints_file = REPO_DIR.joinpath("build", "constraints.txt")
     if not constraints_file.exists():
@@ -33,6 +38,8 @@ def pip_install_requirements(requirements_txt="requirements.txt", continue_on_fa
         constraints_parameters = ["-c", str(constraints_file.resolve())]
     if no_build_isolation:
         constraints_parameters.append("--no-build-isolation")
+    if extra_args and isinstance(extra_args, list):
+        constraints_parameters.extend(extra_args)
     if not continue_on_fail:
         subprocess.check_call(
             [sys.executable, "-m", "pip", "install", "-r", requirements_txt] + constraints_parameters,