Add --python argument to benchmark script (#5844)

## Summary

Makes it easy to benchmark on different Python versions. Used this for
transformers, where some projects don't support Python 3.12.
This commit is contained in:
Charlie Marsh
2024-08-06 21:45:46 -04:00
committed by GitHub
parent cccb6820db
commit 0f63173400
+56 -50
View File
@@ -158,7 +158,8 @@ class Suite(abc.ABC):
class PipCompile(Suite):
def __init__(self, path: str | None = None) -> None:
def __init__(self, *, python: str, path: str | None = None) -> None:
self.python = python
self.name = path or "pip-compile"
self.path = path or "pip-compile"
@@ -282,7 +283,8 @@ class PipCompile(Suite):
class PipSync(Suite):
def __init__(self, path: str | None = None) -> None:
def __init__(self, *, python: str, path: str | None = None) -> None:
self.python = python
self.name = path or "pip-sync"
self.path = path or "pip-sync"
@@ -302,7 +304,7 @@ class PipSync(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_COLD.value})",
prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12.3 {venv_dir}",
prepare=f"rm -rf {cache_dir} && virtualenv --clear -p {self.python} {venv_dir}",
command=[
self.path,
os.path.abspath(requirements_file),
@@ -319,7 +321,7 @@ class PipSync(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_WARM.value})",
prepare=f"virtualenv --clear -p 3.12.3 {venv_dir}",
prepare=f"virtualenv --clear -p {self.python} {venv_dir}",
command=[
self.path,
os.path.abspath(requirements_file),
@@ -332,7 +334,8 @@ class PipSync(Suite):
class Poetry(Suite):
def __init__(self, path: str | None = None) -> None:
def __init__(self, *, python: str, path: str | None = None) -> None:
self.python = python
self.name = path or "poetry"
self.path = path or "poetry"
@@ -359,7 +362,7 @@ class Poetry(Suite):
"bench",
"--no-interaction",
"--python",
"3.12.3",
self.python,
],
cwd=cwd,
stdout=subprocess.DEVNULL,
@@ -559,7 +562,7 @@ class Poetry(Suite):
f"rm -rf {config_dir} && "
f"rm -rf {cache_dir} && "
f"rm -rf {data_dir} &&"
f"virtualenv --clear -p 3.12.3 {venv_dir} --no-seed"
f"virtualenv --clear -p {self.python} {venv_dir} --no-seed"
),
command=[
f"POETRY_CONFIG_DIR={config_dir}",
@@ -599,7 +602,7 @@ class Poetry(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_WARM.value})",
prepare=f"virtualenv --clear -p 3.12.3 {venv_dir}",
prepare=f"virtualenv --clear -p {self.python} {venv_dir}",
command=[
f"POETRY_CONFIG_DIR={config_dir}",
f"POETRY_CACHE_DIR={cache_dir}",
@@ -616,7 +619,8 @@ class Poetry(Suite):
class Pdm(Suite):
def __init__(self, path: str | None = None) -> None:
def __init__(self, *, python: str, path: str | None = None) -> None:
self.python = python
self.name = path or "pdm"
self.path = path or "pdm"
@@ -636,7 +640,7 @@ class Pdm(Suite):
# Create a PDM project.
subprocess.check_call(
[self.path, "init", "--non-interactive", "--python", "3.12.3"],
[self.path, "init", "--non-interactive", "--python", self.python],
cwd=cwd,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
@@ -792,7 +796,7 @@ class Pdm(Suite):
prepare=(
f"rm -rf {cache_dir} && "
f"{self.path} config cache_dir {cache_dir} && "
f"virtualenv --clear -p 3.12.3 {venv_dir} --no-seed"
f"virtualenv --clear -p {self.python} {venv_dir} --no-seed"
),
command=[
f"VIRTUAL_ENV={venv_dir}",
@@ -826,7 +830,7 @@ class Pdm(Suite):
name=f"{self.name} ({Benchmark.INSTALL_WARM.value})",
prepare=(
f"{self.path} config cache_dir {cache_dir} && "
f"virtualenv --clear -p 3.12.3 {venv_dir} --no-seed"
f"virtualenv --clear -p {self.python} {venv_dir} --no-seed"
),
command=[
f"VIRTUAL_ENV={venv_dir}",
@@ -839,8 +843,9 @@ class Pdm(Suite):
class UvPip(Suite):
def __init__(self, *, path: str | None = None) -> Command | None:
def __init__(self, *, python: str, path: str | None = None) -> None:
"""Initialize a uv benchmark."""
self.python = python
self.name = path or "uv pip"
self.path = path or os.path.join(
os.path.dirname(
@@ -967,7 +972,7 @@ class UvPip(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_COLD.value})",
prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12.3 {venv_dir}",
prepare=f"rm -rf {cache_dir} && virtualenv --clear -p {self.python} {venv_dir}",
command=[
f"VIRTUAL_ENV={venv_dir}",
self.path,
@@ -985,7 +990,7 @@ class UvPip(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_WARM.value})",
prepare=f"virtualenv --clear -p 3.12.3 {venv_dir}",
prepare=f"virtualenv --clear -p {self.python} {venv_dir}",
command=[
f"VIRTUAL_ENV={venv_dir}",
self.path,
@@ -999,8 +1004,9 @@ class UvPip(Suite):
class UvProject(Suite):
def __init__(self, *, path: str | None = None) -> Command | None:
def __init__(self, *, python: str, path: str | None = None) -> None:
"""Initialize a uv benchmark."""
self.python = python
self.name = path or "uv"
self.path = path or os.path.join(
os.path.dirname(
@@ -1031,14 +1037,7 @@ class UvProject(Suite):
# Create a Poetry project.
subprocess.check_call(
[
self.path,
"init",
"--name",
"bench",
"--python",
"3.12.3",
],
[self.path, "init", "--name", "bench", "--python", self.python],
cwd=cwd,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
@@ -1073,7 +1072,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1094,7 +1093,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1147,7 +1146,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1180,7 +1179,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1207,7 +1206,7 @@ class UvProject(Suite):
name=f"{self.name} ({Benchmark.INSTALL_COLD.value})",
prepare=(
f"rm -rf {cache_dir} && "
f"virtualenv --clear -p 3.12.3 {venv_dir} --no-seed"
f"virtualenv --clear -p {self.python} {venv_dir} --no-seed"
),
command=[
f"VIRTUAL_ENV={venv_dir}",
@@ -1218,7 +1217,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1243,7 +1242,7 @@ class UvProject(Suite):
return Command(
name=f"{self.name} ({Benchmark.INSTALL_COLD.value})",
prepare=(f"virtualenv --clear -p 3.12.3 {venv_dir} --no-seed"),
prepare=(f"virtualenv --clear -p {self.python} {venv_dir} --no-seed"),
command=[
f"VIRTUAL_ENV={venv_dir}",
self.path,
@@ -1253,7 +1252,7 @@ class UvProject(Suite):
"--directory",
cwd,
"--python",
"3.12.3",
self.python,
],
)
@@ -1275,6 +1274,12 @@ def main():
"--verbose", "-v", action="store_true", help="Print verbose output."
)
parser.add_argument("--json", action="store_true", help="Export results to JSON.")
parser.add_argument(
"--python",
type=str,
help="The Python version to use when benchmarking (e.g., `3.11.7`).",
default="3.12.3",
)
parser.add_argument(
"--warmup",
type=int,
@@ -1376,6 +1381,7 @@ def main():
verbose = args.verbose
json = args.json
python = args.python
warmup = args.warmup
min_runs = args.min_runs
runs = args.runs
@@ -1387,39 +1393,39 @@ def main():
# Determine the tools to benchmark, based on the user-provided arguments.
suites = []
if args.pip_sync:
suites.append(PipSync())
suites.append(PipSync(python=python))
if args.pip_compile:
suites.append(PipCompile())
suites.append(PipCompile(python=python))
if args.poetry:
suites.append(Poetry())
suites.append(Poetry(python=python))
if args.pdm:
suites.append(Pdm())
suites.append(Pdm(python=python))
if args.uv_pip:
suites.append(UvPip())
suites.append(UvPip(python=python))
if args.uv_project:
suites.append(UvProject())
suites.append(UvProject(python=python))
for path in args.pip_sync_path or []:
suites.append(PipSync(path=path))
suites.append(PipSync(python=python, path=path))
for path in args.pip_compile_path or []:
suites.append(PipCompile(path=path))
suites.append(PipCompile(python=python, path=path))
for path in args.poetry_path or []:
suites.append(Poetry(path=path))
suites.append(Poetry(python=python, path=path))
for path in args.pdm_path or []:
suites.append(Pdm(path=path))
suites.append(Pdm(python=python, path=path))
for path in args.uv_pip_path or []:
suites.append(UvPip(path=path))
suites.append(UvPip(python=python, path=path))
for path in args.uv_project_path or []:
suites.append(UvProject(path=path))
suites.append(UvProject(python=python, path=path))
# If no tools were specified, benchmark all tools.
if not suites:
suites = [
PipSync(),
PipCompile(),
Poetry(),
Pdm(),
UvPip(),
UvProject(),
PipSync(python=python),
PipCompile(python=python),
Poetry(python=python),
Pdm(python=python),
UvPip(python=python),
UvProject(python=python),
]
# Determine the benchmarks to run, based on user input. If no benchmarks were