Skip to content

install_requirements.py: use argparse, minor cleanup #7703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jan 17, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
347 changes: 188 additions & 159 deletions install_requirements.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
# LICENSE file in the root directory of this source tree.


import argparse
import glob
import itertools
import os
import platform
import re
Expand Down Expand Up @@ -63,174 +65,201 @@ def python_is_compatible():
return True


if not python_is_compatible():
sys.exit(1)
def clean():
print("Cleaning build artifacts...")
print("Cleaning pip-out/...")
shutil.rmtree("pip-out/", ignore_errors=True)
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
for d in dirs:
print(f"Cleaning {d}...")
shutil.rmtree(d, ignore_errors=True)
print("Done cleaning build artifacts.")

# Parse options.

EXECUTORCH_BUILD_PYBIND = ""
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
USE_PYTORCH_NIGHTLY = True
VALID_PYBINDS = ["coreml", "mps", "xnnpack"]

args = sys.argv[1:]
for arg in args:
if arg == "--pybind":
pass
elif arg in ["coreml", "mps", "xnnpack"]:
if "--pybind" in args:
arg_upper = arg.upper()
EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{arg_upper}=ON"
else:
print(f"Error: {arg} must follow --pybind")
sys.exit(1)
elif arg == "off":
if "--pybind" in args:
if EXECUTORCH_BUILD_PYBIND == "ON":
print("Cannot turnoff pybind option as it is already set.")
sys.exit(1)

def main(args):
if not python_is_compatible():
sys.exit(1)

# Parse options.

EXECUTORCH_BUILD_PYBIND = ""
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
USE_PYTORCH_NIGHTLY = True

parser = argparse.ArgumentParser()
parser.add_argument(
"--pybind",
action="append",
nargs="+",
help="one or more of coreml/mps/xnnpack, or off",
)
parser.add_argument(
"--clean",
action="store_true",
help="clean build artifacts and pip-out instead of installing",
)
parser.add_argument(
"--use-pt-pinned-commit",
action="store_true",
help="build from the pinned PyTorch commit instead of nightly",
)
args = parser.parse_args(args)

if args.clean:
clean()
return

if args.pybind:
# Flatten list of lists.
args.pybind = list(itertools.chain(*args.pybind))
if "off" in args.pybind:
if len(args.pybind) != 1:
raise Exception(
f"Cannot combine `off` with other pybinds: {args.pybind}"
)
EXECUTORCH_BUILD_PYBIND = "OFF"
else:
print(f"Error: {arg} must follow --pybind")
sys.exit(1)

elif arg == "--clean":
print("Cleaning build artifacts...")
print("Cleaning pip-out/...")
shutil.rmtree("pip-out/", ignore_errors=True)
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
for d in dirs:
print(f"Cleaning {d}...")
shutil.rmtree(d, ignore_errors=True)
print("Done cleaning build artifacts.")
sys.exit(0)
elif arg == "--use-pt-pinned-commit":
for pybind_arg in args.pybind:
if pybind_arg not in VALID_PYBINDS:
raise Exception(
f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
)
EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"

if args.use_pt_pinned_commit:
# This option is used in CI to make sure that PyTorch build from the pinned commit
# is used instead of nightly. CI jobs wouldn't be able to catch regression from the
# latest PT commit otherwise
USE_PYTORCH_NIGHTLY = False
else:
print(f"Error: Unknown option {arg}")
sys.exit(1)

# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
# or is not turned off explicitly (--pybind off)
# then install XNNPACK by default.
if EXECUTORCH_BUILD_PYBIND == "":
EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"

# Use ClangCL on Windows.
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
if os.name == "nt":
CMAKE_ARGS += " -T ClangCL"

# Since ExecuTorch often uses main-branch features of pytorch, only the nightly
# pip versions will have the required features.
#
# NOTE: If a newly-fetched version of the executorch repo changes the value of
# NIGHTLY_VERSION, you should re-run this script to install the necessary
# package versions.
NIGHTLY_VERSION = "dev20250104"

# The pip repository that hosts nightly torch packages.
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"

# pip packages needed by exir.
EXIR_REQUIREMENTS = [
# Setting USE_PYTORCH_NIGHTLY to false to test the pinned PyTorch commit. Note
# that we don't need to set any version number there because they have already
# been installed on CI before this step, so pip won't reinstall them
f"torch==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torch",
(
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
if USE_PYTORCH_NIGHTLY
else "torchvision"
), # For testing.
"typing-extensions",
]

# pip packages needed to run examples.
# TODO: Make each example publish its own requirements.txt
EXAMPLES_REQUIREMENTS = [
"timm==1.0.7",
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
"torchsr==1.0.4",
"transformers==4.47.1",
]

# pip packages needed for development.
DEVEL_REQUIREMENTS = [
"cmake", # For building binary targets.
"pip>=23", # For building the pip package.
"pyyaml", # Imported by the kernel codegen tools.
"setuptools>=63", # For building the pip package.
"tomli", # Imported by extract_sources.py when using python < 3.11.
"wheel", # For building the pip package archive.
"zstd", # Imported by resolve_buck.py.
]

# Assemble the list of requirements to actually install.
# TODO: Add options for reducing the number of requirements.
REQUIREMENTS_TO_INSTALL = EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS

# Install the requirements. `--extra-index-url` tells pip to look for package
# versions on the provided URL if they aren't available on the default URL.
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
*REQUIREMENTS_TO_INSTALL,
"--extra-index-url",
TORCH_NIGHTLY_URL,
],
check=True,
)

LOCAL_REQUIREMENTS = [
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
]

# Install packages directly from local copy instead of pypi.
# This is usually not recommended.
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
*LOCAL_REQUIREMENTS,
],
check=True,
)
# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
# or is not turned off explicitly (--pybind off)
# then install XNNPACK by default.
if EXECUTORCH_BUILD_PYBIND == "":
EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"

# Use ClangCL on Windows.
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
if os.name == "nt":
CMAKE_ARGS += " -T ClangCL"

# Since ExecuTorch often uses main-branch features of pytorch, only the nightly
# pip versions will have the required features.
#
# NOTE: If a newly-fetched version of the executorch repo changes the value of
# NIGHTLY_VERSION, you should re-run this script to install the necessary
# package versions.
NIGHTLY_VERSION = "dev20250104"

# The pip repository that hosts nightly torch packages.
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"

# pip packages needed by exir.
EXIR_REQUIREMENTS = [
# Setting USE_PYTORCH_NIGHTLY to false to test the pinned PyTorch commit. Note
# that we don't need to set any version number there because they have already
# been installed on CI before this step, so pip won't reinstall them
f"torch==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torch",
(
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
if USE_PYTORCH_NIGHTLY
else "torchvision"
), # For testing.
"typing-extensions",
]

# pip packages needed to run examples.
# TODO: Make each example publish its own requirements.txt
EXAMPLES_REQUIREMENTS = [
"timm==1.0.7",
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
"torchsr==1.0.4",
"transformers==4.47.1",
]

# pip packages needed for development.
DEVEL_REQUIREMENTS = [
"cmake", # For building binary targets.
"pip>=23", # For building the pip package.
"pyyaml", # Imported by the kernel codegen tools.
"setuptools>=63", # For building the pip package.
"tomli", # Imported by extract_sources.py when using python < 3.11.
"wheel", # For building the pip package archive.
"zstd", # Imported by resolve_buck.py.
]

# Assemble the list of requirements to actually install.
# TODO: Add options for reducing the number of requirements.
REQUIREMENTS_TO_INSTALL = (
EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS
)

# Install the requirements. `--extra-index-url` tells pip to look for package
# versions on the provided URL if they aren't available on the default URL.
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
*REQUIREMENTS_TO_INSTALL,
"--extra-index-url",
TORCH_NIGHTLY_URL,
],
check=True,
)

LOCAL_REQUIREMENTS = [
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
]

# Install packages directly from local copy instead of pypi.
# This is usually not recommended.
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
*LOCAL_REQUIREMENTS,
],
check=True,
)

#
# Install executorch pip package. This also makes `flatc` available on the path.
# The --extra-index-url may be necessary if pyproject.toml has a dependency on a
# pre-release or nightly version of a torch package.
#

# Set environment variables
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS

# Run the pip install command
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
".",
"--no-build-isolation",
"-v",
"--extra-index-url",
TORCH_NIGHTLY_URL,
],
check=True,
)

#
# Install executorch pip package. This also makes `flatc` available on the path.
# The --extra-index-url may be necessary if pyproject.toml has a dependency on a
# pre-release or nightly version of a torch package.
#

# Set environment variables
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS

# Run the pip install command
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
".",
"--no-build-isolation",
"-v",
"--extra-index-url",
TORCH_NIGHTLY_URL,
],
check=True,
)
if __name__ == "__main__":
main(sys.argv[1:])