Allow using alternative pip installers in installation scripts

torchbenchmark/canary_models/DALLE2_pytorch/install.py

@@ -3,7 +3,8 @@
 import sys
 
 import patch
-from utils.python_utils import pip_install_requirements
+
+from utils.python_utils import get_pip_cmd, pip_install_requirements
 
 
 def patch_dalle2():
@@ -24,9 +25,7 @@ def pip_install_requirements_dalle2():
     # embedding-reader requires an old version of pandas and pyarrow
     # https://github.com/rom1504/embedding-reader/blob/a4fd55830a502685600ed8ef07947cd1cb92b083/requirements.txt#L5
     # So we need to reinstall a newer version of pandas and pyarrow, to be compatible with other models
-    subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", "-U", "pandas", "pyarrow"]
-    )
+    subprocess.check_call(get_pip_cmd() + ["install", "-U", "pandas", "pyarrow"])
 
 
 if __name__ == "__main__":

torchbenchmark/canary_models/fambench_xlmr/install.py

@@ -3,7 +3,7 @@
 import sys
 
 from torchbenchmark import REPO_PATH
-from utils.python_utils import pip_install_requirements
+from utils.python_utils import get_pip_cmd, pip_install_requirements
 
 
 def update_fambench_submodule():
@@ -25,10 +25,8 @@ def pip_install_requirements():
         # pin fairseq version
         # ignore deps specified in requirements.txt
         subprocess.check_call(
-            [
-                sys.executable,
-                "-m",
-                "pip",
+            get_pip_cmd
+            + [
                 "install",
                 "--no-deps",
                 "git+https://github.com/facebookresearch/fairseq.git@ae59bd6",

torchbenchmark/canary_models/gcn/install.py

@@ -2,14 +2,13 @@
 import sys
 
 from utils import s3_utils
+from utils.python_utils import get_pip_cmd
 
 
 def pip_install_requirements():
     subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
+        get_pip_cmd()
+        + [
             "install",
             "-q",
             "-r",

torchbenchmark/e2e_models/fambench_xlmr/install.py

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

torchbenchmark/e2e_models/hf_bert/install.py

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

torchbenchmark/e2e_models/hf_t5/install.py

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

torchbenchmark/models/BERT_pytorch/install.py

@@ -1,9 +1,10 @@
 import subprocess
-import sys
+
+from utils.python_utils import get_pip_cmd
 
 
 def setup_install():
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."])
+    subprocess.check_call(get_pip_cmd() + ["install", "-e", "."])
 
 
 if __name__ == "__main__":

torchbenchmark/models/hf_Roberta_base/install.py

@@ -6,11 +6,12 @@
     cache_model,
     patch_transformers,
 )
+from utils.python_utils import get_pip_cmd
 
 
 def pip_install_requirements():
     subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"]
+        get_pip_cmd() + ["install", "-q", "-r", "requirements.txt"]
     )
 
 

torchbenchmark/util/framework/detectron2/__init__.py

@@ -6,7 +6,7 @@
 from urllib import request
 
 from utils import s3_utils
-from utils.python_utils import pip_install_requirements
+from utils.python_utils import get_pip_cmd, pip_install_requirements
 
 CURRENT_DIR = Path(os.path.dirname(os.path.realpath(__file__)))
 # Load pre-trained weights
@@ -46,15 +46,8 @@ def pip_install_requirements_detectron2():
     # Installing by --no-build-isolation after explicitly installing build-time requirements is required.
     # See https://github.com/facebookresearch/detectron2/issues/4921
     subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            "-q",
-            "wheel",
-            "cython",
-        ]  # Build-time requirements
+        get_pip_cmd()
+        + ["install", "-q", "wheel", "cython"]  # Build-time requirements
     )
     pip_install_requirements(
         requirements_txt=requirements_file, no_build_isolation=True

userbenchmark/distributed/install.py

@@ -1,10 +1,11 @@
 import subprocess
-import sys
+
+from utils.python_utils import get_pip_cmd
 
 
 def pip_install_requirements():
     subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"]
+        get_pip_cmd() + ["install", "-q", "-r", "requirements.txt"]
     )
 
 

userbenchmark/dynamo/dynamobench/huggingface.py

@@ -10,7 +10,6 @@
 import types
 import warnings
 
-
 try:
     from .common import (
         BenchmarkRunner,
@@ -32,6 +31,7 @@
 from torch._dynamo.testing import collect_results
 from torch._dynamo.utils import clone_inputs
 
+from utils.python_utils import get_pip_cmd
 
 log = logging.getLogger(__name__)
 
@@ -45,7 +45,7 @@
 
 
 def pip_install(package):
-    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+    subprocess.check_call(get_pip_cmd() + ["install", package])
 
 
 # Disable the flake warnings for the imports. Flake8 does not provide a way to
@@ -647,7 +647,9 @@ def refresh_model_names_and_batch_sizes():
                 + [f"--output={MODELS_FILENAME}"]
             )
         except subprocess.SubprocessError:
-            log.warning(f"Failed to find suitable batch size for {model_name}")  # noqa: G004
+            log.warning(
+                f"Failed to find suitable batch size for {model_name}"
+            )  # noqa: G004
 
 
 def huggingface_main():

userbenchmark/dynamo/dynamobench/timm_models.py

@@ -8,7 +8,6 @@
 import sys
 import warnings
 
-
 try:
     from .common import BenchmarkRunner, download_retry_decorator, load_yaml_file, main
 except ImportError:
@@ -18,14 +17,15 @@
 from torch._dynamo.testing import collect_results, reduce_to_scalar_loss
 from torch._dynamo.utils import clone_inputs
 
+from utils.python_utils import get_pip_cmd
 
 # Enable FX graph caching
 if "TORCHINDUCTOR_FX_GRAPH_CACHE" not in os.environ:
     torch._inductor.config.fx_graph_cache = True
 
 
 def pip_install(package):
-    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+    subprocess.check_call(get_pip_cmd() + ["install", package])
 
 
 try:

userbenchmark/test_bench/install.py

@@ -2,9 +2,10 @@
 import os
 import subprocess
 import sys
-
 from typing import Optional, Tuple
 
+from utils.python_utils import get_pip_cmd
+
 parser = argparse.ArgumentParser(allow_abbrev=False)
 parser.add_argument(
     "models",
@@ -45,7 +46,7 @@ def run_install_py(cwd: str, continue_on_fail=False) -> Tuple[bool, Optional[str
 def pip_install_requirements(
     cwd: str, continue_on_fail=False
 ) -> Tuple[bool, Optional[str]]:
-    cmd = [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"]
+    cmd = get_pip_cmd() + ["install", "-r", "requirements.txt"]
     return _run(*cmd, cwd=cwd, continue_on_fail=continue_on_fail)
 
 

utils/python_utils.py

@@ -1,7 +1,9 @@
+import os
+import shutil
 import subprocess
+import sys
 import warnings
 from pathlib import Path
-
 from typing import List, Optional
 
 DEFAULT_PYTHON_VERSION = "3.12"
@@ -25,14 +27,20 @@ def create_conda_env(pyver: str, name: str):
     subprocess.check_call(command)
 
 
+def get_pip_cmd():
+    if env := os.getenv("PIP_MODULE"):
+        return env.split()
+    else:
+        return [sys.executable, "-m", "pip"]
+
+
 def pip_install_requirements(
     requirements_txt="requirements.txt",
     continue_on_fail=False,
     no_build_isolation=False,
     extra_args: Optional[List[str]] = None,
 ):
-    import sys
-
+    install_cmd = get_pip_cmd()
     constraints_file = REPO_DIR.joinpath("build", "constraints.txt")
     if not constraints_file.exists():
         warnings.warn(
@@ -50,14 +58,12 @@ def pip_install_requirements(
         constraints_parameters.extend(extra_args)
     if not continue_on_fail:
         subprocess.check_call(
-            [sys.executable, "-m", "pip", "install", "-r", requirements_txt]
-            + constraints_parameters,
+            install_cmd + ["install", "-r", requirements_txt] + constraints_parameters,
         )
         return True, None
     try:
         subprocess.run(
-            [sys.executable, "-m", "pip", "install", "-r", requirements_txt]
-            + constraints_parameters,
+            install_cmd + ["install", "-r", requirements_txt] + constraints_parameters,
             check=True,
             stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,