Skip to content

install_requirements.py: refactor: extract install_requirements() function #7704

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 12 commits into from
Jan 17, 2025
Merged

install_requirements.py: refactor: extract install_requirements() function #7704

Changes from 2 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
da1df94
Update
swolchok Jan 16, 2025
e411a1e
Update
swolchok Jan 16, 2025
9a1ddde
Update
swolchok Jan 16, 2025
2063432
Update
swolchok Jan 16, 2025
a893391
Update
swolchok Jan 16, 2025
bf753a8
Update
swolchok Jan 16, 2025
9158d0e
Update
swolchok Jan 16, 2025
e089ac5
Update
swolchok Jan 16, 2025
2dcdbcd
Update
swolchok Jan 16, 2025
1d4dfab
Update
swolchok Jan 16, 2025
ff3a43d
Update
swolchok Jan 16, 2025
641ca89
Update
swolchok Jan 17, 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: 186 additions & 157 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,69 +65,20 @@ 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

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)
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":
# 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"
# The pip repository that hosts nightly torch packages.
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"

# 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.
Expand All @@ -135,103 +88,179 @@ def python_is_compatible():
# 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,
)
def install_requirements(use_pytorch_nightly):
# 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,
)


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(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:
for pybind_arg in args.pybind:
if pybind_arg not in ["coreml", "mps", "xnnpack"]:
continue
EXECUTORCH_BUILD_PYBIND = "ON"
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"

if args.clean:
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

install_requirements(use_pytorch_nightly)

# 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"

#
# 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