diff --git a/torchbenchmark/canary_models/DALLE2_pytorch/install.py b/torchbenchmark/canary_models/DALLE2_pytorch/install.py index f3d509cc1d..cc33cc2cac 100644 --- a/torchbenchmark/canary_models/DALLE2_pytorch/install.py +++ b/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__": diff --git a/torchbenchmark/canary_models/fambench_xlmr/install.py b/torchbenchmark/canary_models/fambench_xlmr/install.py index a09abe998e..55f6bae6a9 100644 --- a/torchbenchmark/canary_models/fambench_xlmr/install.py +++ b/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", diff --git a/torchbenchmark/canary_models/gcn/install.py b/torchbenchmark/canary_models/gcn/install.py index 6f907a3d9f..c9a15111cc 100644 --- a/torchbenchmark/canary_models/gcn/install.py +++ b/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", diff --git a/torchbenchmark/e2e_models/fambench_xlmr/install.py b/torchbenchmark/e2e_models/fambench_xlmr/install.py index e92d8783ca..e57005432e 100644 --- a/torchbenchmark/e2e_models/fambench_xlmr/install.py +++ b/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"] ) diff --git a/torchbenchmark/e2e_models/hf_bert/install.py b/torchbenchmark/e2e_models/hf_bert/install.py index e92d8783ca..e57005432e 100644 --- a/torchbenchmark/e2e_models/hf_bert/install.py +++ b/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"] ) diff --git a/torchbenchmark/e2e_models/hf_t5/install.py b/torchbenchmark/e2e_models/hf_t5/install.py index e92d8783ca..e57005432e 100644 --- a/torchbenchmark/e2e_models/hf_t5/install.py +++ b/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"] ) diff --git a/torchbenchmark/models/BERT_pytorch/install.py b/torchbenchmark/models/BERT_pytorch/install.py index c93fd5c955..7cdc0f6ec7 100644 --- a/torchbenchmark/models/BERT_pytorch/install.py +++ b/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__": diff --git a/torchbenchmark/models/hf_Roberta_base/install.py b/torchbenchmark/models/hf_Roberta_base/install.py index 6c6b31995c..d73448e662 100644 --- a/torchbenchmark/models/hf_Roberta_base/install.py +++ b/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"] ) diff --git a/torchbenchmark/util/framework/detectron2/__init__.py b/torchbenchmark/util/framework/detectron2/__init__.py index c98e649452..ed4319c5ab 100644 --- a/torchbenchmark/util/framework/detectron2/__init__.py +++ b/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 diff --git a/userbenchmark/distributed/install.py b/userbenchmark/distributed/install.py index e92d8783ca..3714866185 100644 --- a/userbenchmark/distributed/install.py +++ b/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"] ) diff --git a/userbenchmark/dynamo/dynamobench/huggingface.py b/userbenchmark/dynamo/dynamobench/huggingface.py index d856a241cc..2bf43a07e2 100755 --- a/userbenchmark/dynamo/dynamobench/huggingface.py +++ b/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(): diff --git a/userbenchmark/dynamo/dynamobench/timm_models.py b/userbenchmark/dynamo/dynamobench/timm_models.py index f31dbb3f95..655e34b4cc 100755 --- a/userbenchmark/dynamo/dynamobench/timm_models.py +++ b/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,6 +17,7 @@ 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: @@ -25,7 +25,7 @@ def pip_install(package): - subprocess.check_call([sys.executable, "-m", "pip", "install", package]) + subprocess.check_call(get_pip_cmd() + ["install", package]) try: diff --git a/userbenchmark/test_bench/install.py b/userbenchmark/test_bench/install.py index f0ad4523e1..0328d87964 100644 --- a/userbenchmark/test_bench/install.py +++ b/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) diff --git a/utils/python_utils.py b/utils/python_utils.py index e1815079c2..31e5d80ad5 100644 --- a/utils/python_utils.py +++ b/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,