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.

Sign up for GitHub

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

install_requirements.py: use argparse, minor cleanup #7703

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
da1df94
Update
swolchok Jan 16, 2025
9a1ddde
Update
swolchok Jan 16, 2025
bf753a8
Update
swolchok Jan 16, 2025
e089ac5
Update
swolchok Jan 16, 2025
1d4dfab
Update
swolchok Jan 16, 2025
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
343 changes: 183 additions & 160 deletions install_requirements.py
View file
Open in desktop
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,175 +65,196 @@ 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
def main(args):
if not python_is_compatible():
sys.exit(1)

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)
# 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(prog="install_requirements")
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.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 ["coreml", "mps", "xnnpack"]:
continue
Copy link
Contributor

@mergennachin mergennachin Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

raise exception here?

Copy link
Contributor Author

@swolchok swolchok Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

whoops! I misinterpreted the original code after I'd made some modifications to it. will fix.

EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"

if args.clean:
Copy link
Contributor

@mergennachin mergennachin Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do this if statement first?

clean()
return

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.
"ai-edge-model-explorer>=0.1.16", # For visualizing ExportedPrograms
]

# 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.
"ai-edge-model-explorer>=0.1.16", # For visualizing ExportedPrograms
]

# 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:])
Loading