From 3626d08cca45aee28adb21da3a0ee38a1fcb521c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 28 Jul 2024 19:27:14 -0400 Subject: [PATCH] Enable benchmarking of `uv tool` and `pipx` (#5531) ## Summary Closes https://github.com/astral-sh/uv/issues/5263. --- BENCHMARKS.md | 4 +- CONTRIBUTING.md | 2 +- scripts/benchmark/README.md | 2 +- scripts/benchmark/pyproject.toml | 4 +- scripts/benchmark/src/benchmark/__init__.py | 1300 +------------------ scripts/benchmark/src/benchmark/resolver.py | 1289 ++++++++++++++++++ scripts/benchmark/src/benchmark/tools.py | 338 +++++ scripts/benchmark/uv.lock | 38 + 8 files changed, 1676 insertions(+), 1301 deletions(-) create mode 100644 scripts/benchmark/src/benchmark/resolver.py create mode 100644 scripts/benchmark/src/benchmark/tools.py diff --git a/BENCHMARKS.md b/BENCHMARKS.md index bf51fde26768..9e6d9125ca6f 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -67,7 +67,7 @@ The benchmark script itself has a several requirements: To benchmark resolution against pip-compile, Poetry, and PDM: ```shell -uv run benchmark \ +uv run resolver \ --uv-pip \ --poetry \ --pdm \ @@ -80,7 +80,7 @@ uv run benchmark \ To benchmark installation against pip-sync, Poetry, and PDM: ```shell -uv run benchmark \ +uv run resolver \ --uv-pip \ --poetry \ --pdm \ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 52d4a7d62958..d807e0225e45 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -86,7 +86,7 @@ We provide diverse sets of requirements for testing and benchmarking the resolve You can use `scripts/benchmark` to benchmark predefined workloads between uv versions and with other tools, e.g., from the `scripts/benchmark` directory: ```shell -uv run benchmark \ +uv run resolver \ --uv-pip \ --poetry \ --benchmark \ diff --git a/scripts/benchmark/README.md b/scripts/benchmark/README.md index d0936ccb8835..1134f7448ea9 100644 --- a/scripts/benchmark/README.md +++ b/scripts/benchmark/README.md @@ -7,7 +7,7 @@ Benchmarking scripts for uv and other package management tools. From the `scripts/benchmark` directory: ```shell -uv run benchmark \ +uv run resolver \ --uv-pip \ --poetry \ --benchmark \ diff --git a/scripts/benchmark/pyproject.toml b/scripts/benchmark/pyproject.toml index e4af335db68e..ac75ddce5bae 100644 --- a/scripts/benchmark/pyproject.toml +++ b/scripts/benchmark/pyproject.toml @@ -10,10 +10,12 @@ dependencies = [ "tomli", "tomli_w", "virtualenv", + "pipx", ] [project.scripts] -benchmark = "benchmark:main" +resolver = "benchmark.resolver:main" +tools = "benchmark.tools:main" [build-system] requires = ["hatchling"] diff --git a/scripts/benchmark/src/benchmark/__init__.py b/scripts/benchmark/src/benchmark/__init__.py index b0e80ef71735..1fd5fa299e48 100644 --- a/scripts/benchmark/src/benchmark/__init__.py +++ b/scripts/benchmark/src/benchmark/__init__.py @@ -1,87 +1,13 @@ -"""Benchmark uv against other packaging tools. - -This script assumes that Python 3.12 is installed. - -By default, this script also assumes that `pip`, `pip-tools`, `virtualenv`, `poetry` and -`hyperfine` are installed, and that a uv release builds exists at `./target/release/uv` -(relative to the repository root). However, the set of tools is configurable. - -For example, to benchmark uv's `pip compile` command against `pip-tools`, run the -following from the `scripts/benchmark` directory: - - uv run benchmark --uv-pip --pip-compile ../requirements/trio.in - -It's most common to benchmark multiple uv versions against one another by building -from multiple branches and specifying the path to each binary, as in: - - # Build the baseline version, from the repo root. - git checkout main - cargo build --release - mv ./target/release/uv ./target/release/baseline - - # Build the feature version, again from the repo root. - git checkout feature - cargo build --release - - # Run the benchmark. - cd scripts/benchmark - uv run benchmark \ - --uv-pip-path ../../target/release/uv \ - --uv-pip-path ../../target/release/baseline \ - ../requirements/trio.in - -By default, the script will run the resolution benchmarks when a `requirements.in` file -is provided, and the installation benchmarks when a `requirements.txt` file is provided: - - # Run the resolution benchmarks against the Trio project. - uv run bench \ - --uv-path ../../target/release/uv \ - --uv-path ../../target/release/baseline \ - ../requirements/trio.in - - # Run the installation benchmarks against the Trio project. - uv run bench \ - --uv-path ../../target/release/uv \ - --uv-path ../../target/release/baseline \ - ../requirements/compiled/trio.txt - -You can also specify the benchmark to run explicitly: - - # Run the "uncached install" benchmark against the Trio project. - uv run bench \ - --uv-path ../../target/release/uv \ - --uv-path ../../target/release/baseline \ - --benchmark install-cold \ - ../requirements/compiled/trio.txt -""" - -import abc -import argparse -import enum -import logging -import os.path import shlex -import shutil import subprocess -import tempfile import typing -class Benchmark(enum.Enum): - """Enumeration of the benchmarks to run.""" - - RESOLVE_COLD = "resolve-cold" - RESOLVE_WARM = "resolve-warm" - RESOLVE_INCREMENTAL = "resolve-incremental" - INSTALL_COLD = "install-cold" - INSTALL_WARM = "install-warm" - - class Command(typing.NamedTuple): name: str """The name of the command to benchmark.""" - prepare: str + prepare: str | None """The command to run before each benchmark run.""" command: list[str] @@ -89,7 +15,7 @@ class Command(typing.NamedTuple): class Hyperfine(typing.NamedTuple): - benchmark: Benchmark + name: str """The benchmark to run.""" commands: list[Command] @@ -114,7 +40,7 @@ def run(self) -> None: # Export to JSON. if self.json: args.append("--export-json") - args.append(f"{self.benchmark.value}.json") + args.append(f"{self.name}.json") # Preamble: benchmark-wide setup. if self.verbose: @@ -132,1228 +58,10 @@ def run(self) -> None: # Add all prepare statements. for command in self.commands: args.append("--prepare") - args.append(command.prepare) + args.append(command.prepare or "") # Add all commands. for command in self.commands: args.append(shlex.join(command.command)) subprocess.check_call(args) - - -# The requirement to use when benchmarking an incremental resolution. -# Ideally, this requirement is compatible with all requirements files, but does not -# appear in any resolutions. -INCREMENTAL_REQUIREMENT = "django" - - -class Suite(abc.ABC): - """Abstract base class for packaging tools.""" - - def command( - self, - benchmark: Benchmark, - requirements_file: str, - *, - cwd: str, - ) -> Command | None: - """Generate a command to benchmark a given tool.""" - match benchmark: - case Benchmark.RESOLVE_COLD: - return self.resolve_cold(requirements_file, cwd=cwd) - case Benchmark.RESOLVE_WARM: - return self.resolve_warm(requirements_file, cwd=cwd) - case Benchmark.RESOLVE_INCREMENTAL: - return self.resolve_incremental(requirements_file, cwd=cwd) - case Benchmark.INSTALL_COLD: - return self.install_cold(requirements_file, cwd=cwd) - case Benchmark.INSTALL_WARM: - return self.install_warm(requirements_file, cwd=cwd) - case _: - raise ValueError(f"Invalid benchmark: {benchmark}") - - @abc.abstractmethod - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - """Resolve a set of dependencies using pip-tools, from a cold cache. - - The resolution is performed from scratch, i.e., without an existing lockfile, - and the cache directory is cleared between runs. - """ - - @abc.abstractmethod - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - """Resolve a set of dependencies using pip-tools, from a warm cache. - - The resolution is performed from scratch, i.e., without an existing lockfile; - however, the cache directory is _not_ cleared between runs. - """ - - @abc.abstractmethod - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - """Resolve a modified lockfile using pip-tools, from a warm cache. - - The resolution is performed with an existing lockfile, and the cache directory - is _not_ cleared between runs. However, a new dependency is added to the set - of input requirements, which does not appear in the lockfile. - """ - - @abc.abstractmethod - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - """Install a set of dependencies using pip-tools, from a cold cache. - - The virtual environment is recreated before each run, and the cache directory - is cleared between runs. - """ - - @abc.abstractmethod - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - """Install a set of dependencies using pip-tools, from a cold cache. - - The virtual environment is recreated before each run, and the cache directory - is cleared between runs. - """ - - -class PipCompile(Suite): - def __init__(self, path: str | None = None) -> None: - self.name = path or "pip-compile" - self.path = path or "pip-compile" - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "requirements.txt") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", - prepare=f"rm -rf {cwd} && rm -f {output_file}", - command=[ - self.path, - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - output_file, - "--rebuild", - ], - ) - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "requirements.txt") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", - prepare=f"rm -f {output_file}", - command=[ - self.path, - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - output_file, - ], - ) - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - baseline = os.path.join(cwd, "baseline.txt") - - # First, perform a cold resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [ - self.path, - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - baseline, - ], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(baseline), f"Lock file doesn't exist at: {baseline}" - - input_file = os.path.join(cwd, "requirements.in") - output_file = os.path.join(cwd, "requirements.txt") - - # Add a dependency to the requirements file. - with open(input_file, "w") as fp1: - fp1.write(f"{INCREMENTAL_REQUIREMENT}\n") - with open(requirements_file) as fp2: - fp1.writelines(fp2.readlines()) - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", - prepare=f"rm -f {output_file} && cp {baseline} {output_file}", - command=[ - self.path, - input_file, - "--cache-dir", - cache_dir, - "--output-file", - output_file, - ], - ) - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: ... - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: ... - - -class PipSync(Suite): - def __init__(self, path: str | None = None) -> None: - self.name = path or "pip-sync" - self.path = path or "pip-sync" - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: ... - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: ... - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: ... - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12 {venv_dir}", - command=[ - self.path, - os.path.abspath(requirements_file), - "--pip-args", - f"--cache-dir {cache_dir}", - "--python-executable", - os.path.join(venv_dir, "bin", "python"), - ], - ) - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", - prepare=f"virtualenv --clear -p 3.12 {venv_dir}", - command=[ - self.path, - os.path.abspath(requirements_file), - "--pip-args", - f"--cache-dir {cache_dir}", - "--python-executable", - os.path.join(venv_dir, "bin", "python"), - ], - ) - - -class Poetry(Suite): - def __init__(self, path: str | None = None) -> None: - self.name = path or "poetry" - self.path = path or "poetry" - - def setup(self, requirements_file: str, *, cwd: str) -> None: - """Initialize a Poetry project from a requirements file.""" - import tomli - import tomli_w - from packaging.requirements import Requirement - - # Parse all dependencies from the requirements file. - with open(requirements_file) as fp: - requirements = [ - Requirement(line) - for line in fp - if not line.lstrip().startswith("#") and len(line.strip()) > 0 - ] - - # Create a Poetry project. - subprocess.check_call( - [ - self.path, - "init", - "--name", - "bench", - "--no-interaction", - "--python", - "3.12.3", - ], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - # Parse the pyproject.toml. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["tool"]["poetry"]["dependencies"].update( - { - str(requirement.name): str(requirement.specifier) - if requirement.specifier - else "*" - for requirement in requirements - } - ) - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - poetry_lock = os.path.join(cwd, "poetry.lock") - config_dir = os.path.join(cwd, "config", "pypoetry") - cache_dir = os.path.join(cwd, "cache", "pypoetry") - data_dir = os.path.join(cwd, "data", "pypoetry") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", - prepare=( - f"rm -rf {config_dir} && " - f"rm -rf {cache_dir} && " - f"rm -rf {data_dir} &&" - f"rm -rf {poetry_lock}" - ), - command=[ - f"POETRY_CONFIG_DIR={config_dir}", - f"POETRY_CACHE_DIR={cache_dir}", - f"POETRY_DATA_DIR={data_dir}", - self.path, - "lock", - "--directory", - cwd, - ], - ) - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - poetry_lock = os.path.join(cwd, "poetry.lock") - config_dir = os.path.join(cwd, "config", "pypoetry") - cache_dir = os.path.join(cwd, "cache", "pypoetry") - data_dir = os.path.join(cwd, "data", "pypoetry") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", - prepare=f"rm -f {poetry_lock}", - command=[ - f"POETRY_CONFIG_DIR={config_dir}", - f"POETRY_CACHE_DIR={cache_dir}", - f"POETRY_DATA_DIR={data_dir}", - self.path, - "lock", - "--directory", - cwd, - ], - ) - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - import tomli - import tomli_w - - self.setup(requirements_file, cwd=cwd) - - poetry_lock = os.path.join(cwd, "poetry.lock") - assert not os.path.exists( - poetry_lock - ), f"Lock file already exists at: {poetry_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" - - # Add a dependency to the requirements file. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["tool"]["poetry"]["dependencies"].update( - { - INCREMENTAL_REQUIREMENT: "*", - } - ) - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - # Store the baseline lockfile. - baseline = os.path.join(cwd, "baseline.lock") - shutil.copyfile(poetry_lock, baseline) - - poetry_lock = os.path.join(cwd, "poetry.lock") - config_dir = os.path.join(cwd, "config", "pypoetry") - cache_dir = os.path.join(cwd, "cache", "pypoetry") - data_dir = os.path.join(cwd, "data", "pypoetry") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", - prepare=f"rm {poetry_lock} && cp {baseline} {poetry_lock}", - command=[ - f"POETRY_CONFIG_DIR={config_dir}", - f"POETRY_CACHE_DIR={cache_dir}", - f"POETRY_DATA_DIR={data_dir}", - self.path, - "lock", - "--no-update", - "--directory", - cwd, - ], - ) - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - poetry_lock = os.path.join(cwd, "poetry.lock") - assert not os.path.exists( - poetry_lock - ), f"Lock file already exists at: {poetry_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" - - config_dir = os.path.join(cwd, "config", "pypoetry") - cache_dir = os.path.join(cwd, "cache", "pypoetry") - data_dir = os.path.join(cwd, "data", "pypoetry") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=( - f"rm -rf {config_dir} && " - f"rm -rf {cache_dir} && " - f"rm -rf {data_dir} &&" - f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" - ), - command=[ - f"POETRY_CONFIG_DIR={config_dir}", - f"POETRY_CACHE_DIR={cache_dir}", - f"POETRY_DATA_DIR={data_dir}", - f"VIRTUAL_ENV={venv_dir}", - self.path, - "install", - "--no-root", - "--sync", - "--directory", - cwd, - ], - ) - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - poetry_lock = os.path.join(cwd, "poetry.lock") - assert not os.path.exists( - poetry_lock - ), f"Lock file already exists at: {poetry_lock}" - - # Run a resolution, to ensure that the lockfile exists. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" - - config_dir = os.path.join(cwd, "config", "pypoetry") - cache_dir = os.path.join(cwd, "cache", "pypoetry") - data_dir = os.path.join(cwd, "data", "pypoetry") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", - prepare=f"virtualenv --clear -p 3.12 {venv_dir}", - command=[ - f"POETRY_CONFIG_DIR={config_dir}", - f"POETRY_CACHE_DIR={cache_dir}", - f"POETRY_DATA_DIR={data_dir}", - f"VIRTUAL_ENV={venv_dir}", - self.path, - "install", - "--no-root", - "--sync", - "--directory", - cwd, - ], - ) - - -class Pdm(Suite): - def __init__(self, path: str | None = None) -> None: - self.name = path or "pdm" - self.path = path or "pdm" - - def setup(self, requirements_file: str, *, cwd: str) -> None: - """Initialize a PDM project from a requirements file.""" - import tomli - import tomli_w - from packaging.requirements import Requirement - - # Parse all dependencies from the requirements file. - with open(requirements_file) as fp: - requirements = [ - Requirement(line) - for line in fp - if not line.lstrip().startswith("#") and len(line.strip()) > 0 - ] - - # Create a PDM project. - subprocess.check_call( - [self.path, "init", "--non-interactive", "--python", "3.12"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - # Parse the pyproject.toml. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["project"]["dependencies"] = [ - str(requirement) for requirement in requirements - ] - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - pdm_lock = os.path.join(cwd, "pdm.lock") - cache_dir = os.path.join(cwd, "cache", "pdm") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", - prepare=f"rm -rf {cache_dir} && rm -rf {pdm_lock} && {self.path} config cache_dir {cache_dir}", - command=[ - self.path, - "lock", - "--project", - cwd, - ], - ) - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - pdm_lock = os.path.join(cwd, "pdm.lock") - cache_dir = os.path.join(cwd, "cache", "pdm") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", - prepare=f"rm -rf {pdm_lock} && {self.path} config cache_dir {cache_dir}", - command=[ - self.path, - "lock", - "--project", - cwd, - ], - ) - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - import tomli - import tomli_w - - self.setup(requirements_file, cwd=cwd) - - pdm_lock = os.path.join(cwd, "pdm.lock") - assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" - - # Add a dependency to the requirements file. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["project"]["dependencies"] += [INCREMENTAL_REQUIREMENT] - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - # Store the baseline lockfile. - baseline = os.path.join(cwd, "baseline.lock") - shutil.copyfile(pdm_lock, baseline) - - pdm_lock = os.path.join(cwd, "pdm.lock") - cache_dir = os.path.join(cwd, "cache", "pdm") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", - prepare=f"rm -f {pdm_lock} && cp {baseline} {pdm_lock} && {self.path} config cache_dir {cache_dir}", - command=[ - self.path, - "lock", - "--update-reuse", - "--project", - cwd, - ], - ) - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - pdm_lock = os.path.join(cwd, "pdm.lock") - assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" - - venv_dir = os.path.join(cwd, ".venv") - cache_dir = os.path.join(cwd, "cache", "pdm") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=( - f"rm -rf {cache_dir} && " - f"{self.path} config cache_dir {cache_dir} && " - f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" - ), - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "sync", - "--project", - cwd, - ], - ) - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - pdm_lock = os.path.join(cwd, "pdm.lock") - assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" - - venv_dir = os.path.join(cwd, ".venv") - cache_dir = os.path.join(cwd, "cache", "pdm") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", - prepare=( - f"{self.path} config cache_dir {cache_dir} && " - f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" - ), - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "sync", - "--project", - cwd, - ], - ) - - -class UvPip(Suite): - def __init__(self, *, path: str | None = None) -> Command | None: - """Initialize a uv benchmark.""" - self.name = path or "uv" - self.path = path or os.path.join( - os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ), - "target", - "release", - "uv", - ) - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "requirements.txt") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", - prepare=f"rm -rf {cache_dir} && rm -f {output_file}", - command=[ - self.path, - "pip", - "compile", - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - output_file, - ], - ) - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "requirements.txt") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", - prepare=f"rm -f {output_file}", - command=[ - self.path, - "pip", - "compile", - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - output_file, - ], - ) - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - baseline = os.path.join(cwd, "baseline.txt") - - # First, perform a cold resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [ - self.path, - "pip", - "compile", - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - "--output-file", - baseline, - ], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(baseline), f"Lock file doesn't exist at: {baseline}" - - input_file = os.path.join(cwd, "requirements.in") - output_file = os.path.join(cwd, "requirements.txt") - - # Add a dependency to the requirements file. - with open(input_file, "w") as fp1: - fp1.write(f"{INCREMENTAL_REQUIREMENT}\n") - with open(requirements_file) as fp2: - fp1.writelines(fp2.readlines()) - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", - prepare=f"rm -f {output_file} && cp {baseline} {output_file}", - command=[ - self.path, - "pip", - "compile", - input_file, - "--cache-dir", - cache_dir, - "--output-file", - output_file, - ], - ) - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12 {venv_dir}", - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "pip", - "sync", - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - ], - ) - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", - prepare=f"virtualenv --clear -p 3.12 {venv_dir}", - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "pip", - "sync", - os.path.abspath(requirements_file), - "--cache-dir", - cache_dir, - ], - ) - - -class UvProject(Suite): - def __init__(self, *, path: str | None = None) -> Command | None: - """Initialize a uv benchmark.""" - self.name = path or "uv" - self.path = path or os.path.join( - os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ), - "target", - "release", - "uv", - ) - - def setup(self, requirements_file: str, *, cwd: str) -> None: - """Initialize a uv project from a requirements file.""" - import tomli - import tomli_w - from packaging.requirements import Requirement - - # Parse all dependencies from the requirements file. - with open(requirements_file) as fp: - requirements = [ - Requirement(line) - for line in fp - if not line.lstrip().startswith("#") and len(line.strip()) > 0 - ] - - # Create a Poetry project. - subprocess.check_call( - [ - self.path, - "init", - "--name", - "bench", - "--python", - "3.12.3", - ], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - # Parse the pyproject.toml. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["project"]["dependencies"] += [ - str(requirement) for requirement in requirements - ] - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "uv.lock") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", - prepare=f"rm -rf {cache_dir} && rm -f {output_file}", - command=[ - self.path, - "lock", - "--cache-dir", - cache_dir, - "--directory", - cwd, - "--python", - "3.12.3", - ], - ) - - def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - cache_dir = os.path.join(cwd, ".cache") - output_file = os.path.join(cwd, "uv.lock") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", - prepare=f"rm -f {output_file}", - command=[ - self.path, - "lock", - "--cache-dir", - cache_dir, - "--directory", - cwd, - "--python", - "3.12.3", - ], - ) - - def resolve_incremental( - self, requirements_file: str, *, cwd: str - ) -> Command | None: - import tomli - import tomli_w - - self.setup(requirements_file, cwd=cwd) - - uv_lock = os.path.join(cwd, "uv.lock") - assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" - - # Add a dependency to the requirements file. - with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: - pyproject = tomli.load(fp) - - # Add the dependencies to the pyproject.toml. - pyproject["project"]["dependencies"] += [INCREMENTAL_REQUIREMENT] - - with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: - tomli_w.dump(pyproject, fp) - - # Store the baseline lockfile. - baseline = os.path.join(cwd, "baseline.lock") - shutil.copyfile(uv_lock, baseline) - - uv_lock = os.path.join(cwd, "uv.lock") - cache_dir = os.path.join(cwd, ".cache") - - return Command( - name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", - prepare=f"rm -f {uv_lock} && cp {baseline} {uv_lock}", - command=[ - self.path, - "lock", - "--cache-dir", - cache_dir, - "--directory", - cwd, - "--python", - "3.12.3", - ], - ) - - def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - uv_lock = os.path.join(cwd, "uv.lock") - assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" - - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=( - f"rm -rf {cache_dir} && " - f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" - ), - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "sync", - "--cache-dir", - cache_dir, - "--directory", - cwd, - "--python", - "3.12.3", - ], - ) - - def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: - self.setup(requirements_file, cwd=cwd) - - uv_lock = os.path.join(cwd, "uv.lock") - assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" - - # Run a resolution, to ensure that the lockfile exists. - # TODO(charlie): Make this a `setup`. - subprocess.check_call( - [self.path, "lock"], - cwd=cwd, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" - - cache_dir = os.path.join(cwd, ".cache") - venv_dir = os.path.join(cwd, ".venv") - - return Command( - name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", - prepare=(f"virtualenv --clear -p 3.12 {venv_dir} --no-seed"), - command=[ - f"VIRTUAL_ENV={venv_dir}", - self.path, - "sync", - "--cache-dir", - cache_dir, - "--directory", - cwd, - "--python", - "3.12.3", - ], - ) - - -def main(): - """Run the benchmark.""" - parser = argparse.ArgumentParser( - description="Benchmark uv against other packaging tools." - ) - parser.add_argument( - "file", - type=str, - help=( - "The file to read the dependencies from (typically: `requirements.in` " - "(for resolution) or `requirements.txt` (for installation))." - ), - ) - parser.add_argument( - "--verbose", "-v", action="store_true", help="Print verbose output." - ) - parser.add_argument("--json", action="store_true", help="Export results to JSON.") - parser.add_argument( - "--warmup", - type=int, - help="The number of warmup runs to perform.", - default=3, - ) - parser.add_argument( - "--min-runs", - type=int, - help="The minimum number of runs to perform.", - default=10, - ) - parser.add_argument( - "--benchmark", - "-b", - type=str, - help="The benchmark(s) to run.", - choices=[benchmark.value for benchmark in Benchmark], - action="append", - ) - parser.add_argument( - "--pip-sync", - help="Whether to benchmark `pip-sync` (requires `pip-tools` to be installed).", - action="store_true", - ) - parser.add_argument( - "--pip-compile", - help="Whether to benchmark `pip-compile` (requires `pip-tools` to be installed).", - action="store_true", - ) - parser.add_argument( - "--poetry", - help="Whether to benchmark Poetry (requires Poetry to be installed).", - action="store_true", - ) - parser.add_argument( - "--pdm", - help="Whether to benchmark PDM (requires PDM to be installed).", - action="store_true", - ) - parser.add_argument( - "--uv-pip", - help="Whether to benchmark uv's pip interface (assumes a uv binary exists at `./target/release/uv`).", - action="store_true", - ) - parser.add_argument( - "--uv-project", - help="Whether to benchmark uv's project interface (assumes a uv binary exists at `./target/release/uv`).", - action="store_true", - ) - parser.add_argument( - "--pip-sync-path", - type=str, - help="Path(s) to the `pip-sync` binary to benchmark.", - action="append", - ) - parser.add_argument( - "--pip-compile-path", - type=str, - help="Path(s) to the `pip-compile` binary to benchmark.", - action="append", - ) - parser.add_argument( - "--poetry-path", - type=str, - help="Path(s) to the Poetry binary to benchmark.", - action="append", - ) - parser.add_argument( - "--pdm-path", - type=str, - help="Path(s) to the PDM binary to benchmark.", - action="append", - ) - parser.add_argument( - "--uv-pip-path", - type=str, - help="Path(s) to the uv binary to benchmark.", - action="append", - ) - parser.add_argument( - "--uv-project-path", - type=str, - help="Path(s) to the uv binary to benchmark.", - action="append", - ) - - args = parser.parse_args() - logging.basicConfig( - level=logging.INFO if args.verbose else logging.WARN, - format="%(asctime)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - - verbose = args.verbose - json = args.json - warmup = args.warmup - min_runs = args.min_runs - - requirements_file = os.path.abspath(args.file) - if not os.path.exists(requirements_file): - raise ValueError(f"File not found: {requirements_file}") - - # Determine the tools to benchmark, based on the user-provided arguments. - suites = [] - if args.pip_sync: - suites.append(PipSync()) - if args.pip_compile: - suites.append(PipCompile()) - if args.poetry: - suites.append(Poetry()) - if args.pdm: - suites.append(Pdm()) - if args.uv_pip: - suites.append(UvPip()) - if args.uv_project: - suites.append(UvProject()) - for path in args.pip_sync_path or []: - suites.append(PipSync(path=path)) - for path in args.pip_compile_path or []: - suites.append(PipCompile(path=path)) - for path in args.poetry_path or []: - suites.append(Poetry(path=path)) - for path in args.pdm_path or []: - suites.append(Pdm(path=path)) - for path in args.uv_pip_path or []: - suites.append(UvPip(path=path)) - for path in args.uv_project_path or []: - suites.append(UvProject(path=path)) - - # If no tools were specified, benchmark all tools. - if not suites: - suites = [ - PipSync(), - PipCompile(), - Poetry(), - UvPip(), - UvProject(), - ] - - # Determine the benchmarks to run, based on user input. If no benchmarks were - # specified, infer an appropriate set based on the file extension. - benchmarks = ( - [Benchmark(benchmark) for benchmark in args.benchmark] - if args.benchmark is not None - else [Benchmark.RESOLVE_COLD, Benchmark.RESOLVE_WARM] - if requirements_file.endswith(".in") - else [Benchmark.INSTALL_COLD, Benchmark.INSTALL_WARM] - if requirements_file.endswith(".txt") - else list(Benchmark) - ) - - logging.info(f"Reading requirements from: {requirements_file}") - logging.info("```") - with open(args.file) as f: - for line in f: - logging.info(line.rstrip()) - logging.info("```") - - with tempfile.TemporaryDirectory() as cwd: - for benchmark in benchmarks: - # Generate the benchmark command for each tool. - commands = [ - command - for suite in suites - if ( - command := suite.command( - benchmark, requirements_file, cwd=tempfile.mkdtemp(dir=cwd) - ) - ) - ] - - if commands: - hyperfine = Hyperfine( - benchmark=benchmark, - commands=commands, - warmup=warmup, - min_runs=min_runs, - verbose=verbose, - json=json, - ) - hyperfine.run() - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark/src/benchmark/resolver.py b/scripts/benchmark/src/benchmark/resolver.py new file mode 100644 index 000000000000..8f38d599b535 --- /dev/null +++ b/scripts/benchmark/src/benchmark/resolver.py @@ -0,0 +1,1289 @@ +"""Benchmark uv against other packaging tools. + +For example, to benchmark uv's `pip compile` command against `pip-tools`, run the +following from the `scripts/benchmark` directory: + + uv run resolver --uv-pip --pip-compile ../requirements/trio.in + +It's most common to benchmark multiple uv versions against one another by building +from multiple branches and specifying the path to each binary, as in: + + # Build the baseline version, from the repo root. + git checkout main + cargo build --release + mv ./target/release/uv ./target/release/baseline + + # Build the feature version, again from the repo root. + git checkout feature + cargo build --release + + # Run the benchmark. + cd scripts/benchmark + uv run resolver \ + --uv-pip-path ../../target/release/uv \ + --uv-pip-path ../../target/release/baseline \ + ../requirements/trio.in + +By default, the script will run the resolution benchmarks when a `requirements.in` file +is provided, and the installation benchmarks when a `requirements.txt` file is provided: + + # Run the resolution benchmarks against the Trio project. + uv run benchmark\ + --uv-path ../../target/release/uv \ + --uv-path ../../target/release/baseline \ + ../requirements/trio.in + + # Run the installation benchmarks against the Trio project. + uv run benchmark\ + --uv-path ../../target/release/uv \ + --uv-path ../../target/release/baseline \ + ../requirements/compiled/trio.txt + +You can also specify the benchmark to run explicitly: + + # Run the "uncached install" benchmark against the Trio project. + uv run benchmark\ + --uv-path ../../target/release/uv \ + --uv-path ../../target/release/baseline \ + --benchmark install-cold \ + ../requirements/compiled/trio.txt +""" + +import abc +import argparse +import enum +import logging +import os.path +import shutil +import subprocess +import tempfile + +from benchmark import Command, Hyperfine + + +class Benchmark(enum.Enum): + """Enumeration of the benchmarks to run.""" + + RESOLVE_COLD = "resolve-cold" + RESOLVE_WARM = "resolve-warm" + RESOLVE_INCREMENTAL = "resolve-incremental" + INSTALL_COLD = "install-cold" + INSTALL_WARM = "install-warm" + + +# The requirement to use when benchmarking an incremental resolution. +# Ideally, this requirement is compatible with all requirements files, but does not +# appear in any resolutions. +INCREMENTAL_REQUIREMENT = "django" + + +class Suite(abc.ABC): + """Abstract base class for packaging tools.""" + + def command( + self, + benchmark: Benchmark, + requirements_file: str, + *, + cwd: str, + ) -> Command | None: + """Generate a command to benchmark a given tool.""" + match benchmark: + case Benchmark.RESOLVE_COLD: + return self.resolve_cold(requirements_file, cwd=cwd) + case Benchmark.RESOLVE_WARM: + return self.resolve_warm(requirements_file, cwd=cwd) + case Benchmark.RESOLVE_INCREMENTAL: + return self.resolve_incremental(requirements_file, cwd=cwd) + case Benchmark.INSTALL_COLD: + return self.install_cold(requirements_file, cwd=cwd) + case Benchmark.INSTALL_WARM: + return self.install_warm(requirements_file, cwd=cwd) + case _: + raise ValueError(f"Invalid benchmark: {benchmark}") + + @abc.abstractmethod + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + """Resolve a set of dependencies using pip-tools, from a cold cache. + + The resolution is performed from scratch, i.e., without an existing lockfile, + and the cache directory is cleared between runs. + """ + + @abc.abstractmethod + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + """Resolve a set of dependencies using pip-tools, from a warm cache. + + The resolution is performed from scratch, i.e., without an existing lockfile; + however, the cache directory is _not_ cleared between runs. + """ + + @abc.abstractmethod + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + """Resolve a modified lockfile using pip-tools, from a warm cache. + + The resolution is performed with an existing lockfile, and the cache directory + is _not_ cleared between runs. However, a new dependency is added to the set + of input requirements, which does not appear in the lockfile. + """ + + @abc.abstractmethod + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + """Install a set of dependencies using pip-tools, from a cold cache. + + The virtual environment is recreated before each run, and the cache directory + is cleared between runs. + """ + + @abc.abstractmethod + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + """Install a set of dependencies using pip-tools, from a cold cache. + + The virtual environment is recreated before each run, and the cache directory + is cleared between runs. + """ + + +class PipCompile(Suite): + def __init__(self, path: str | None = None) -> None: + self.name = path or "pip-compile" + self.path = path or "pip-compile" + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "requirements.txt") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", + prepare=f"rm -rf {cwd} && rm -f {output_file}", + command=[ + self.path, + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + output_file, + "--rebuild", + ], + ) + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "requirements.txt") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", + prepare=f"rm -f {output_file}", + command=[ + self.path, + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + output_file, + ], + ) + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + baseline = os.path.join(cwd, "baseline.txt") + + # First, perform a cold resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [ + self.path, + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + baseline, + ], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(baseline), f"Lock file doesn't exist at: {baseline}" + + input_file = os.path.join(cwd, "requirements.in") + output_file = os.path.join(cwd, "requirements.txt") + + # Add a dependency to the requirements file. + with open(input_file, "w") as fp1: + fp1.write(f"{INCREMENTAL_REQUIREMENT}\n") + with open(requirements_file) as fp2: + fp1.writelines(fp2.readlines()) + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", + prepare=f"rm -f {output_file} && cp {baseline} {output_file}", + command=[ + self.path, + input_file, + "--cache-dir", + cache_dir, + "--output-file", + output_file, + ], + ) + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: ... + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: ... + + +class PipSync(Suite): + def __init__(self, path: str | None = None) -> None: + self.name = path or "pip-sync" + self.path = path or "pip-sync" + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: ... + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: ... + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: ... + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12 {venv_dir}", + command=[ + self.path, + os.path.abspath(requirements_file), + "--pip-args", + f"--cache-dir {cache_dir}", + "--python-executable", + os.path.join(venv_dir, "bin", "python"), + ], + ) + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=f"virtualenv --clear -p 3.12 {venv_dir}", + command=[ + self.path, + os.path.abspath(requirements_file), + "--pip-args", + f"--cache-dir {cache_dir}", + "--python-executable", + os.path.join(venv_dir, "bin", "python"), + ], + ) + + +class Poetry(Suite): + def __init__(self, path: str | None = None) -> None: + self.name = path or "poetry" + self.path = path or "poetry" + + def setup(self, requirements_file: str, *, cwd: str) -> None: + """Initialize a Poetry project from a requirements file.""" + import tomli + import tomli_w + from packaging.requirements import Requirement + + # Parse all dependencies from the requirements file. + with open(requirements_file) as fp: + requirements = [ + Requirement(line) + for line in fp + if not line.lstrip().startswith("#") and len(line.strip()) > 0 + ] + + # Create a Poetry project. + subprocess.check_call( + [ + self.path, + "init", + "--name", + "bench", + "--no-interaction", + "--python", + "3.12.3", + ], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + # Parse the pyproject.toml. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["tool"]["poetry"]["dependencies"].update( + { + str(requirement.name): str(requirement.specifier) + if requirement.specifier + else "*" + for requirement in requirements + } + ) + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + poetry_lock = os.path.join(cwd, "poetry.lock") + config_dir = os.path.join(cwd, "config", "pypoetry") + cache_dir = os.path.join(cwd, "cache", "pypoetry") + data_dir = os.path.join(cwd, "data", "pypoetry") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", + prepare=( + f"rm -rf {config_dir} && " + f"rm -rf {cache_dir} && " + f"rm -rf {data_dir} &&" + f"rm -rf {poetry_lock}" + ), + command=[ + f"POETRY_CONFIG_DIR={config_dir}", + f"POETRY_CACHE_DIR={cache_dir}", + f"POETRY_DATA_DIR={data_dir}", + self.path, + "lock", + "--directory", + cwd, + ], + ) + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + poetry_lock = os.path.join(cwd, "poetry.lock") + config_dir = os.path.join(cwd, "config", "pypoetry") + cache_dir = os.path.join(cwd, "cache", "pypoetry") + data_dir = os.path.join(cwd, "data", "pypoetry") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", + prepare=f"rm -f {poetry_lock}", + command=[ + f"POETRY_CONFIG_DIR={config_dir}", + f"POETRY_CACHE_DIR={cache_dir}", + f"POETRY_DATA_DIR={data_dir}", + self.path, + "lock", + "--directory", + cwd, + ], + ) + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + import tomli + import tomli_w + + self.setup(requirements_file, cwd=cwd) + + poetry_lock = os.path.join(cwd, "poetry.lock") + assert not os.path.exists( + poetry_lock + ), f"Lock file already exists at: {poetry_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" + + # Add a dependency to the requirements file. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["tool"]["poetry"]["dependencies"].update( + { + INCREMENTAL_REQUIREMENT: "*", + } + ) + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + # Store the baseline lockfile. + baseline = os.path.join(cwd, "baseline.lock") + shutil.copyfile(poetry_lock, baseline) + + poetry_lock = os.path.join(cwd, "poetry.lock") + config_dir = os.path.join(cwd, "config", "pypoetry") + cache_dir = os.path.join(cwd, "cache", "pypoetry") + data_dir = os.path.join(cwd, "data", "pypoetry") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", + prepare=f"rm {poetry_lock} && cp {baseline} {poetry_lock}", + command=[ + f"POETRY_CONFIG_DIR={config_dir}", + f"POETRY_CACHE_DIR={cache_dir}", + f"POETRY_DATA_DIR={data_dir}", + self.path, + "lock", + "--no-update", + "--directory", + cwd, + ], + ) + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + poetry_lock = os.path.join(cwd, "poetry.lock") + assert not os.path.exists( + poetry_lock + ), f"Lock file already exists at: {poetry_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" + + config_dir = os.path.join(cwd, "config", "pypoetry") + cache_dir = os.path.join(cwd, "cache", "pypoetry") + data_dir = os.path.join(cwd, "data", "pypoetry") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=( + f"rm -rf {config_dir} && " + f"rm -rf {cache_dir} && " + f"rm -rf {data_dir} &&" + f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" + ), + command=[ + f"POETRY_CONFIG_DIR={config_dir}", + f"POETRY_CACHE_DIR={cache_dir}", + f"POETRY_DATA_DIR={data_dir}", + f"VIRTUAL_ENV={venv_dir}", + self.path, + "install", + "--no-root", + "--sync", + "--directory", + cwd, + ], + ) + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + poetry_lock = os.path.join(cwd, "poetry.lock") + assert not os.path.exists( + poetry_lock + ), f"Lock file already exists at: {poetry_lock}" + + # Run a resolution, to ensure that the lockfile exists. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(poetry_lock), f"Lock file doesn't exist at: {poetry_lock}" + + config_dir = os.path.join(cwd, "config", "pypoetry") + cache_dir = os.path.join(cwd, "cache", "pypoetry") + data_dir = os.path.join(cwd, "data", "pypoetry") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=f"virtualenv --clear -p 3.12 {venv_dir}", + command=[ + f"POETRY_CONFIG_DIR={config_dir}", + f"POETRY_CACHE_DIR={cache_dir}", + f"POETRY_DATA_DIR={data_dir}", + f"VIRTUAL_ENV={venv_dir}", + self.path, + "install", + "--no-root", + "--sync", + "--directory", + cwd, + ], + ) + + +class Pdm(Suite): + def __init__(self, path: str | None = None) -> None: + self.name = path or "pdm" + self.path = path or "pdm" + + def setup(self, requirements_file: str, *, cwd: str) -> None: + """Initialize a PDM project from a requirements file.""" + import tomli + import tomli_w + from packaging.requirements import Requirement + + # Parse all dependencies from the requirements file. + with open(requirements_file) as fp: + requirements = [ + Requirement(line) + for line in fp + if not line.lstrip().startswith("#") and len(line.strip()) > 0 + ] + + # Create a PDM project. + subprocess.check_call( + [self.path, "init", "--non-interactive", "--python", "3.12"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + # Parse the pyproject.toml. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["project"]["dependencies"] = [ + str(requirement) for requirement in requirements + ] + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + pdm_lock = os.path.join(cwd, "pdm.lock") + cache_dir = os.path.join(cwd, "cache", "pdm") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", + prepare=f"rm -rf {cache_dir} && rm -rf {pdm_lock} && {self.path} config cache_dir {cache_dir}", + command=[ + self.path, + "lock", + "--project", + cwd, + ], + ) + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + pdm_lock = os.path.join(cwd, "pdm.lock") + cache_dir = os.path.join(cwd, "cache", "pdm") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", + prepare=f"rm -rf {pdm_lock} && {self.path} config cache_dir {cache_dir}", + command=[ + self.path, + "lock", + "--project", + cwd, + ], + ) + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + import tomli + import tomli_w + + self.setup(requirements_file, cwd=cwd) + + pdm_lock = os.path.join(cwd, "pdm.lock") + assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" + + # Add a dependency to the requirements file. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["project"]["dependencies"] += [INCREMENTAL_REQUIREMENT] + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + # Store the baseline lockfile. + baseline = os.path.join(cwd, "baseline.lock") + shutil.copyfile(pdm_lock, baseline) + + pdm_lock = os.path.join(cwd, "pdm.lock") + cache_dir = os.path.join(cwd, "cache", "pdm") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", + prepare=f"rm -f {pdm_lock} && cp {baseline} {pdm_lock} && {self.path} config cache_dir {cache_dir}", + command=[ + self.path, + "lock", + "--update-reuse", + "--project", + cwd, + ], + ) + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + pdm_lock = os.path.join(cwd, "pdm.lock") + assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" + + venv_dir = os.path.join(cwd, ".venv") + cache_dir = os.path.join(cwd, "cache", "pdm") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=( + f"rm -rf {cache_dir} && " + f"{self.path} config cache_dir {cache_dir} && " + f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" + ), + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "sync", + "--project", + cwd, + ], + ) + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + pdm_lock = os.path.join(cwd, "pdm.lock") + assert not os.path.exists(pdm_lock), f"Lock file already exists at: {pdm_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(pdm_lock), f"Lock file doesn't exist at: {pdm_lock}" + + venv_dir = os.path.join(cwd, ".venv") + cache_dir = os.path.join(cwd, "cache", "pdm") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=( + f"{self.path} config cache_dir {cache_dir} && " + f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" + ), + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "sync", + "--project", + cwd, + ], + ) + + +class UvPip(Suite): + def __init__(self, *, path: str | None = None) -> Command | None: + """Initialize a uv benchmark.""" + self.name = path or "uv" + self.path = path or os.path.join( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ), + "target", + "release", + "uv", + ) + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "requirements.txt") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", + prepare=f"rm -rf {cache_dir} && rm -f {output_file}", + command=[ + self.path, + "pip", + "compile", + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + output_file, + ], + ) + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "requirements.txt") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", + prepare=f"rm -f {output_file}", + command=[ + self.path, + "pip", + "compile", + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + output_file, + ], + ) + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + baseline = os.path.join(cwd, "baseline.txt") + + # First, perform a cold resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [ + self.path, + "pip", + "compile", + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + "--output-file", + baseline, + ], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(baseline), f"Lock file doesn't exist at: {baseline}" + + input_file = os.path.join(cwd, "requirements.in") + output_file = os.path.join(cwd, "requirements.txt") + + # Add a dependency to the requirements file. + with open(input_file, "w") as fp1: + fp1.write(f"{INCREMENTAL_REQUIREMENT}\n") + with open(requirements_file) as fp2: + fp1.writelines(fp2.readlines()) + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", + prepare=f"rm -f {output_file} && cp {baseline} {output_file}", + command=[ + self.path, + "pip", + "compile", + input_file, + "--cache-dir", + cache_dir, + "--output-file", + output_file, + ], + ) + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=f"rm -rf {cache_dir} && virtualenv --clear -p 3.12 {venv_dir}", + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "pip", + "sync", + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + ], + ) + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=f"virtualenv --clear -p 3.12 {venv_dir}", + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "pip", + "sync", + os.path.abspath(requirements_file), + "--cache-dir", + cache_dir, + ], + ) + + +class UvProject(Suite): + def __init__(self, *, path: str | None = None) -> Command | None: + """Initialize a uv benchmark.""" + self.name = path or "uv" + self.path = path or os.path.join( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ), + "target", + "release", + "uv", + ) + + def setup(self, requirements_file: str, *, cwd: str) -> None: + """Initialize a uv project from a requirements file.""" + import tomli + import tomli_w + from packaging.requirements import Requirement + + # Parse all dependencies from the requirements file. + with open(requirements_file) as fp: + requirements = [ + Requirement(line) + for line in fp + if not line.lstrip().startswith("#") and len(line.strip()) > 0 + ] + + # Create a Poetry project. + subprocess.check_call( + [ + self.path, + "init", + "--name", + "bench", + "--python", + "3.12.3", + ], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + # Parse the pyproject.toml. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["project"]["dependencies"] += [ + str(requirement) for requirement in requirements + ] + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + def resolve_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "uv.lock") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_COLD.value})", + prepare=f"rm -rf {cache_dir} && rm -f {output_file}", + command=[ + self.path, + "lock", + "--cache-dir", + cache_dir, + "--directory", + cwd, + "--python", + "3.12.3", + ], + ) + + def resolve_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + cache_dir = os.path.join(cwd, ".cache") + output_file = os.path.join(cwd, "uv.lock") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_WARM.value})", + prepare=f"rm -f {output_file}", + command=[ + self.path, + "lock", + "--cache-dir", + cache_dir, + "--directory", + cwd, + "--python", + "3.12.3", + ], + ) + + def resolve_incremental( + self, requirements_file: str, *, cwd: str + ) -> Command | None: + import tomli + import tomli_w + + self.setup(requirements_file, cwd=cwd) + + uv_lock = os.path.join(cwd, "uv.lock") + assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" + + # Add a dependency to the requirements file. + with open(os.path.join(cwd, "pyproject.toml"), "rb") as fp: + pyproject = tomli.load(fp) + + # Add the dependencies to the pyproject.toml. + pyproject["project"]["dependencies"] += [INCREMENTAL_REQUIREMENT] + + with open(os.path.join(cwd, "pyproject.toml"), "wb") as fp: + tomli_w.dump(pyproject, fp) + + # Store the baseline lockfile. + baseline = os.path.join(cwd, "baseline.lock") + shutil.copyfile(uv_lock, baseline) + + uv_lock = os.path.join(cwd, "uv.lock") + cache_dir = os.path.join(cwd, ".cache") + + return Command( + name=f"{self.name} ({Benchmark.RESOLVE_INCREMENTAL.value})", + prepare=f"rm -f {uv_lock} && cp {baseline} {uv_lock}", + command=[ + self.path, + "lock", + "--cache-dir", + cache_dir, + "--directory", + cwd, + "--python", + "3.12.3", + ], + ) + + def install_cold(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + uv_lock = os.path.join(cwd, "uv.lock") + assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" + + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=( + f"rm -rf {cache_dir} && " + f"virtualenv --clear -p 3.12 {venv_dir} --no-seed" + ), + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "sync", + "--cache-dir", + cache_dir, + "--directory", + cwd, + "--python", + "3.12.3", + ], + ) + + def install_warm(self, requirements_file: str, *, cwd: str) -> Command | None: + self.setup(requirements_file, cwd=cwd) + + uv_lock = os.path.join(cwd, "uv.lock") + assert not os.path.exists(uv_lock), f"Lock file already exists at: {uv_lock}" + + # Run a resolution, to ensure that the lockfile exists. + # TODO(charlie): Make this a `setup`. + subprocess.check_call( + [self.path, "lock"], + cwd=cwd, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert os.path.exists(uv_lock), f"Lock file doesn't exist at: {uv_lock}" + + cache_dir = os.path.join(cwd, ".cache") + venv_dir = os.path.join(cwd, ".venv") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=(f"virtualenv --clear -p 3.12 {venv_dir} --no-seed"), + command=[ + f"VIRTUAL_ENV={venv_dir}", + self.path, + "sync", + "--cache-dir", + cache_dir, + "--directory", + cwd, + "--python", + "3.12.3", + ], + ) + + +def main(): + """Run the benchmark.""" + parser = argparse.ArgumentParser( + description="Benchmark uv against other packaging tools." + ) + parser.add_argument( + "file", + type=str, + help=( + "The file to read the dependencies from (typically: `requirements.in` " + "(for resolution) or `requirements.txt` (for installation))." + ), + ) + parser.add_argument( + "--verbose", "-v", action="store_true", help="Print verbose output." + ) + parser.add_argument("--json", action="store_true", help="Export results to JSON.") + parser.add_argument( + "--warmup", + type=int, + help="The number of warmup runs to perform.", + default=3, + ) + parser.add_argument( + "--min-runs", + type=int, + help="The minimum number of runs to perform.", + default=10, + ) + parser.add_argument( + "--benchmark", + "-b", + type=str, + help="The benchmark(s) to run.", + choices=[benchmark.value for benchmark in Benchmark], + action="append", + ) + parser.add_argument( + "--pip-sync", + help="Whether to benchmark `pip-sync` (requires `pip-tools` to be installed).", + action="store_true", + ) + parser.add_argument( + "--pip-compile", + help="Whether to benchmark `pip-compile` (requires `pip-tools` to be installed).", + action="store_true", + ) + parser.add_argument( + "--poetry", + help="Whether to benchmark Poetry (requires Poetry to be installed).", + action="store_true", + ) + parser.add_argument( + "--pdm", + help="Whether to benchmark PDM (requires PDM to be installed).", + action="store_true", + ) + parser.add_argument( + "--uv-pip", + help="Whether to benchmark uv's pip interface (assumes a uv binary exists at `./target/release/uv`).", + action="store_true", + ) + parser.add_argument( + "--uv-project", + help="Whether to benchmark uv's project interface (assumes a uv binary exists at `./target/release/uv`).", + action="store_true", + ) + parser.add_argument( + "--pip-sync-path", + type=str, + help="Path(s) to the `pip-sync` binary to benchmark.", + action="append", + ) + parser.add_argument( + "--pip-compile-path", + type=str, + help="Path(s) to the `pip-compile` binary to benchmark.", + action="append", + ) + parser.add_argument( + "--poetry-path", + type=str, + help="Path(s) to the Poetry binary to benchmark.", + action="append", + ) + parser.add_argument( + "--pdm-path", + type=str, + help="Path(s) to the PDM binary to benchmark.", + action="append", + ) + parser.add_argument( + "--uv-pip-path", + type=str, + help="Path(s) to the uv binary to benchmark.", + action="append", + ) + parser.add_argument( + "--uv-project-path", + type=str, + help="Path(s) to the uv binary to benchmark.", + action="append", + ) + + args = parser.parse_args() + logging.basicConfig( + level=logging.INFO if args.verbose else logging.WARN, + format="%(asctime)s %(levelname)s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + verbose = args.verbose + json = args.json + warmup = args.warmup + min_runs = args.min_runs + + requirements_file = os.path.abspath(args.file) + if not os.path.exists(requirements_file): + raise ValueError(f"File not found: {requirements_file}") + + # Determine the tools to benchmark, based on the user-provided arguments. + suites = [] + if args.pip_sync: + suites.append(PipSync()) + if args.pip_compile: + suites.append(PipCompile()) + if args.poetry: + suites.append(Poetry()) + if args.pdm: + suites.append(Pdm()) + if args.uv_pip: + suites.append(UvPip()) + if args.uv_project: + suites.append(UvProject()) + for path in args.pip_sync_path or []: + suites.append(PipSync(path=path)) + for path in args.pip_compile_path or []: + suites.append(PipCompile(path=path)) + for path in args.poetry_path or []: + suites.append(Poetry(path=path)) + for path in args.pdm_path or []: + suites.append(Pdm(path=path)) + for path in args.uv_pip_path or []: + suites.append(UvPip(path=path)) + for path in args.uv_project_path or []: + suites.append(UvProject(path=path)) + + # If no tools were specified, benchmark all tools. + if not suites: + suites = [ + PipSync(), + PipCompile(), + Poetry(), + UvPip(), + UvProject(), + ] + + # Determine the benchmarks to run, based on user input. If no benchmarks were + # specified, infer an appropriate set based on the file extension. + benchmarks = ( + [Benchmark(benchmark) for benchmark in args.benchmark] + if args.benchmark is not None + else [Benchmark.RESOLVE_COLD, Benchmark.RESOLVE_WARM] + if requirements_file.endswith(".in") + else [Benchmark.INSTALL_COLD, Benchmark.INSTALL_WARM] + if requirements_file.endswith(".txt") + else list(Benchmark) + ) + + logging.info(f"Reading requirements from: {requirements_file}") + logging.info("```") + with open(args.file) as f: + for line in f: + logging.info(line.rstrip()) + logging.info("```") + + with tempfile.TemporaryDirectory() as cwd: + for benchmark in benchmarks: + # Generate the benchmark command for each tool. + commands = [ + command + for suite in suites + if ( + command := suite.command( + benchmark, requirements_file, cwd=tempfile.mkdtemp(dir=cwd) + ) + ) + ] + + if commands: + hyperfine = Hyperfine( + name=str(benchmark.value), + commands=commands, + warmup=warmup, + min_runs=min_runs, + verbose=verbose, + json=json, + ) + hyperfine.run() + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/src/benchmark/tools.py b/scripts/benchmark/src/benchmark/tools.py new file mode 100644 index 000000000000..a421b01a1b63 --- /dev/null +++ b/scripts/benchmark/src/benchmark/tools.py @@ -0,0 +1,338 @@ +"""Benchmark the uv `tool` interface against other packaging tools. + +For example, to benchmark uv against pipx, run the following from the +`scripts/benchmark` directory: + + uv run tools --uv --pipx +""" + +import abc +import argparse +import enum +import logging +import os.path +import tempfile + +from benchmark import Command, Hyperfine + +TOOL = "flask" + + +class Benchmark(enum.Enum): + """Enumeration of the benchmarks to run.""" + + INSTALL_COLD = "install-cold" + INSTALL_WARM = "install-warm" + RUN = "run" + + +class Suite(abc.ABC): + """Abstract base class for packaging tools.""" + + def command(self, benchmark: Benchmark, *, cwd: str) -> Command | None: + """Generate a command to benchmark a given tool.""" + match benchmark: + case Benchmark.INSTALL_COLD: + return self.install_cold(cwd=cwd) + case Benchmark.INSTALL_WARM: + return self.install_warm(cwd=cwd) + case Benchmark.RUN: + return self.run(cwd=cwd) + case _: + raise ValueError(f"Invalid benchmark: {benchmark}") + + @abc.abstractmethod + def install_cold(self, *, cwd: str) -> Command | None: + """Resolve a set of dependencies using pip-tools, from a cold cache. + + The resolution is performed from scratch, i.e., without an existing lockfile, + and the cache directory is cleared between runs. + """ + + @abc.abstractmethod + def install_warm(self, *, cwd: str) -> Command | None: + """Resolve a set of dependencies using pip-tools, from a warm cache. + + The resolution is performed from scratch, i.e., without an existing lockfile; + however, the cache directory is _not_ cleared between runs. + """ + + @abc.abstractmethod + def run(self, *, cwd: str) -> Command | None: + """Resolve a modified lockfile using pip-tools, from a warm cache. + + The resolution is performed with an existing lockfile, and the cache directory + is _not_ cleared between runs. However, a new dependency is added to the set + of input requirements, which does not appear in the lockfile. + """ + + +class Pipx(Suite): + def __init__(self, path: str | None = None) -> None: + self.name = path or "pipx" + self.path = path or "pipx" + + def install_cold(self, *, cwd: str) -> Command | None: + home_dir = os.path.join(cwd, "home") + bin_dir = os.path.join(cwd, "bin") + man_dir = os.path.join(cwd, "man") + + # pipx uses a shared virtualenv directory in `${PIPX_HOME}/shared`, which + # contains pip. If we remove `${PIPX_HOME}/shared`, we're simulating the _first_ + # pipx invocation on a machine, rather than `pipx run` with a cold cache. So, + # instead, we only remove the installed tools, rather than the shared + # dependencies. + venvs_dir = os.path.join(home_dir, "venvs") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=f"rm -rf {venvs_dir} && rm -rf {bin_dir} && rm -rf {man_dir}", + command=[ + f"PIPX_HOME={home_dir}", + f"PIPX_BIN_DIR={bin_dir}", + f"PIPX_MAN_DIR={man_dir}", + self.path, + "install", + "--pip-args=--no-cache-dir", + TOOL, + ], + ) + + def install_warm(self, *, cwd: str) -> Command | None: + home_dir = os.path.join(cwd, "home") + bin_dir = os.path.join(cwd, "bin") + man_dir = os.path.join(cwd, "man") + + # pipx uses a shared virtualenv directory in `${PIPX_HOME}/shared`, which + # contains pip. If we remove `${PIPX_HOME}/shared`, we're simulating the _first_ + # pipx invocation on a machine, rather than `pipx run` with a cold cache. So, + # instead, we only remove the installed tools, rather than the shared + # dependencies. + venvs_dir = os.path.join(home_dir, "venvs") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=f"rm -rf {venvs_dir} && rm -rf {bin_dir} && rm -rf {man_dir}", + command=[ + f"PIPX_HOME={home_dir}", + f"PIPX_BIN_DIR={bin_dir}", + f"PIPX_MAN_DIR={man_dir}", + self.path, + "install", + TOOL, + ], + ) + + def run(self, *, cwd: str) -> Command | None: + home_dir = os.path.join(cwd, "home") + bin_dir = os.path.join(cwd, "bin") + man_dir = os.path.join(cwd, "man") + + return Command( + name=f"{self.name} ({Benchmark.RUN.value})", + prepare="", + command=[ + f"PIPX_HOME={home_dir}", + f"PIPX_BIN_DIR={bin_dir}", + f"PIPX_MAN_DIR={man_dir}", + self.path, + "install", + TOOL, + ], + ) + + +class Uv(Suite): + def __init__(self, *, path: str | None = None) -> Command | None: + """Initialize a uv benchmark.""" + self.name = path or "uv" + self.path = path or os.path.join( + os.path.dirname( + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ) + ), + "target", + "release", + "uv", + ) + + def install_cold(self, *, cwd: str) -> Command | None: + bin_dir = os.path.join(cwd, "bin") + tool_dir = os.path.join(cwd, "tool") + cache_dir = os.path.join(cwd, ".cache") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_COLD.value})", + prepare=f"rm -rf {bin_dir} && rm -rf {tool_dir} && rm -rf {cache_dir}", + command=[ + f"XDG_BIN_HOME={bin_dir}", + f"UV_TOOL_DIR={tool_dir}", + self.path, + "tool", + "install", + "--cache-dir", + cache_dir, + "--", + TOOL, + ], + ) + + def install_warm(self, *, cwd: str) -> Command | None: + bin_dir = os.path.join(cwd, "bin") + tool_dir = os.path.join(cwd, "tool") + cache_dir = os.path.join(cwd, ".cache") + + return Command( + name=f"{self.name} ({Benchmark.INSTALL_WARM.value})", + prepare=f"rm -rf {bin_dir} && rm -rf {tool_dir}", + command=[ + f"XDG_BIN_HOME={bin_dir}", + f"UV_TOOL_DIR={tool_dir}", + self.path, + "tool", + "install", + "--cache-dir", + cache_dir, + "--", + TOOL, + ], + ) + + def run(self, *, cwd: str) -> Command | None: + bin_dir = os.path.join(cwd, "bin") + tool_dir = os.path.join(cwd, "tool") + cache_dir = os.path.join(cwd, ".cache") + + return Command( + name=f"{self.name} ({Benchmark.RUN.value})", + prepare="", + command=[ + f"XDG_BIN_HOME={bin_dir}", + f"UV_TOOL_DIR={tool_dir}", + self.path, + "tool", + "run", + "--cache-dir", + cache_dir, + "--", + TOOL, + "--version", + ], + ) + + +def main(): + """Run the benchmark.""" + parser = argparse.ArgumentParser( + description="Benchmark uv against other packaging tools." + ) + parser.add_argument( + "--verbose", "-v", action="store_true", help="Print verbose output." + ) + parser.add_argument("--json", action="store_true", help="Export results to JSON.") + parser.add_argument( + "--warmup", + type=int, + help="The number of warmup runs to perform.", + default=3, + ) + parser.add_argument( + "--min-runs", + type=int, + help="The minimum number of runs to perform.", + default=10, + ) + parser.add_argument( + "--benchmark", + "-b", + type=str, + help="The benchmark(s) to run.", + choices=[benchmark.value for benchmark in Benchmark], + action="append", + ) + parser.add_argument( + "--pipx", + help="Whether to benchmark `pipx`.", + action="store_true", + ) + parser.add_argument( + "--uv", + help="Whether to benchmark uv (assumes a uv binary exists at `./target/release/uv`).", + action="store_true", + ) + parser.add_argument( + "--pipx-path", + type=str, + help="Path(s) to the `pipx` binary to benchmark.", + action="append", + ) + parser.add_argument( + "--uv-path", + type=str, + help="Path(s) to the uv binary to benchmark.", + action="append", + ) + + args = parser.parse_args() + logging.basicConfig( + level=logging.INFO if args.verbose else logging.WARN, + format="%(asctime)s %(levelname)s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + verbose = args.verbose + json = args.json + warmup = args.warmup + min_runs = args.min_runs + + # Determine the tools to benchmark, based on the user-provided arguments. + suites = [] + if args.pipx: + suites.append(Pipx()) + if args.uv: + suites.append(Uv()) + for path in args.pipx_path or []: + suites.append(Pipx(path=path)) + for path in args.uv_path or []: + suites.append(Uv(path=path)) + + # If no tools were specified, benchmark all tools. + if not suites: + suites = [ + Pipx(), + Uv(), + ] + + # Determine the benchmarks to run, based on user input. + benchmarks = ( + [Benchmark(benchmark) for benchmark in args.benchmark] + if args.benchmark is not None + else list(Benchmark) + ) + + with tempfile.TemporaryDirectory() as cwd: + for benchmark in benchmarks: + # Generate the benchmark command for each tool. + commands = [ + command + for suite in suites + if (command := suite.command(benchmark, cwd=tempfile.mkdtemp(dir=cwd))) + ] + + if commands: + hyperfine = Hyperfine( + name=str(benchmark.value), + commands=commands, + warmup=warmup, + min_runs=min_runs, + verbose=verbose, + json=json, + ) + hyperfine.run() + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/uv.lock b/scripts/benchmark/uv.lock index c5b3f64c4205..c37c01d3b03c 100644 --- a/scripts/benchmark/uv.lock +++ b/scripts/benchmark/uv.lock @@ -14,6 +14,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/a2/10639a79341f6c019dedc95bd48a4928eed9f1d1197f4c04f546fc7ae0ff/anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7", size = 86780 }, ] +[[distribution]] +name = "argcomplete" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/db/ca/45176b8362eb06b68f946c2bf1184b92fc98d739a3f8c790999a257db91f/argcomplete-3.4.0.tar.gz", hash = "sha256:c2abcdfe1be8ace47ba777d4fce319eb13bf8ad9dace8d085dcad6eded88057f", size = 82275 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/29/cba741f3abc1700dda883c4a1dd83f4ae89e4e8654067929d89143df2c58/argcomplete-3.4.0-py3-none-any.whl", hash = "sha256:69a79e083a716173e5532e0fa3bef45f793f4e61096cf52b5a42c0211c8b8aa5", size = 42641 }, +] + [[distribution]] name = "benchmark" version = "0.0.1" @@ -21,6 +30,7 @@ source = { editable = "." } dependencies = [ { name = "pdm" }, { name = "pip-tools" }, + { name = "pipx" }, { name = "poetry" }, { name = "tomli" }, { name = "tomli-w" }, @@ -511,6 +521,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/dc/38f4ce065e92c66f058ea7a368a9c5de4e702272b479c0992059f7693941/pip_tools-7.4.1-py3-none-any.whl", hash = "sha256:4c690e5fbae2f21e87843e89c26191f0d9454f362d8acdbd695716493ec8b3a9", size = 61235 }, ] +[[distribution]] +name = "pipx" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "argcomplete" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "userpath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2d/f3/c04c5cd0a5795fe6bb09d56c4892384e53cb75813fc08e5cbfa4d080664a/pipx-1.6.0.tar.gz", hash = "sha256:840610e00103e3d49ae24b6b51804b60988851a5dd65468adb71e5a97e2699b2", size = 307321 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/2f/ec2454e6784168837dfe0ffad5080c1c2bdf37ee052999cabfd849a56338/pipx-1.6.0-py3-none-any.whl", hash = "sha256:760889dc3aeed7bf4024973bf22ca0c2a891003f52389159ab5cb0c57d9ebff4", size = 77756 }, +] + [[distribution]] name = "pkginfo" version = "1.11.1" @@ -835,6 +861,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ca/1c/89ffc63a9605b583d5df2be791a27bc1a42b7c32bab68d3c8f2f73a98cd4/urllib3-2.2.2-py3-none-any.whl", hash = "sha256:a448b2f64d686155468037e1ace9f2d2199776e17f0a46610480d311f73e3472", size = 121444 }, ] +[[distribution]] +name = "userpath" +version = "1.9.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/b7/30753098208505d7ff9be5b3a32112fb8a4cb3ddfccbbb7ba9973f2e29ff/userpath-1.9.2.tar.gz", hash = "sha256:6c52288dab069257cc831846d15d48133522455d4677ee69a9781f11dbefd815", size = 11140 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/99/3ec6335ded5b88c2f7ed25c56ffd952546f7ed007ffb1e1539dc3b57015a/userpath-1.9.2-py3-none-any.whl", hash = "sha256:2cbf01a23d655a1ff8fc166dfb78da1b641d1ceabf0fe5f970767d380b14e89d", size = 9065 }, +] + [[distribution]] name = "virtualenv" version = "20.26.3"