Skip to content

Commit 98f8ade

Browse files
swolchokswolchok
committed
install_requirements.py: use argparse, minor cleanup
Replace the bespoke argument parser with argparse in preparation for separating requirements installation from building and installing ExecuTorch itself. Test Plan: (with debug print of args.pybind to save time) install_requirements.sh --pybind xnnpack --pybind mps --pybind coreml install_requirements.sh --pybind xnnpack mps coreml install_requirements.sh --clean CI should cover use-pt-pinned-commit ghstack-source-id: a15e696 ghstack-comment-id: 2596312894 Pull Request resolved: #7703
1 parent a727b55 commit 98f8ade

File tree

1 file changed

+183
-160
lines changed

1 file changed

+183
-160
lines changed

‎install_requirements.py

Lines changed: 183 additions & 160 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
# LICENSE file in the root directory of this source tree.
77

88

9+
import argparse
910
import glob
11+
import itertools
1012
import os
1113
import platform
1214
import re
@@ -63,175 +65,196 @@ def python_is_compatible():
6365
return True
6466

6567

66-
if not python_is_compatible():
67-
sys.exit(1)
68+
def clean():
69+
print("Cleaning build artifacts...")
70+
print("Cleaning pip-out/...")
71+
shutil.rmtree("pip-out/", ignore_errors=True)
72+
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
73+
for d in dirs:
74+
print(f"Cleaning {d}...")
75+
shutil.rmtree(d, ignore_errors=True)
76+
print("Done cleaning build artifacts.")
6877

69-
# Parse options.
7078

71-
EXECUTORCH_BUILD_PYBIND = ""
72-
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
73-
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
74-
USE_PYTORCH_NIGHTLY = True
79+
def main(args):
80+
if not python_is_compatible():
81+
sys.exit(1)
7582

76-
args = sys.argv[1:]
77-
for arg in args:
78-
if arg == "--pybind":
79-
pass
80-
elif arg in ["coreml", "mps", "xnnpack"]:
81-
if "--pybind" in args:
82-
arg_upper = arg.upper()
83-
EXECUTORCH_BUILD_PYBIND = "ON"
84-
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{arg_upper}=ON"
85-
else:
86-
print(f"Error: {arg} must follow --pybind")
87-
sys.exit(1)
88-
elif arg == "off":
89-
if "--pybind" in args:
90-
if EXECUTORCH_BUILD_PYBIND == "ON":
91-
print("Cannot turnoff pybind option as it is already set.")
92-
sys.exit(1)
83+
# Parse options.
84+
85+
EXECUTORCH_BUILD_PYBIND = ""
86+
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
87+
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
88+
USE_PYTORCH_NIGHTLY = True
89+
90+
parser = argparse.ArgumentParser(prog="install_requirements")
91+
parser.add_argument(
92+
"--pybind",
93+
action="append",
94+
nargs="+",
95+
help="one or more of coreml/mps/xnnpack, or off",
96+
)
97+
parser.add_argument(
98+
"--clean",
99+
action="store_true",
100+
help="clean build artifacts and pip-out instead of installing",
101+
)
102+
parser.add_argument(
103+
"--use-pt-pinned-commit",
104+
action="store_true",
105+
help="build from the pinned PyTorch commit instead of nightly",
106+
)
107+
args = parser.parse_args(args)
108+
if args.pybind:
109+
# Flatten list of lists.
110+
args.pybind = list(itertools.chain(*args.pybind))
111+
if "off" in args.pybind:
112+
if len(args.pybind) != 1:
113+
raise Exception(
114+
f"Cannot combine `off` with other pybinds: {args.pybind}"
115+
)
93116
EXECUTORCH_BUILD_PYBIND = "OFF"
94117
else:
95-
print(f"Error: {arg} must follow --pybind")
96-
sys.exit(1)
97-
98-
elif arg == "--clean":
99-
print("Cleaning build artifacts...")
100-
print("Cleaning pip-out/...")
101-
shutil.rmtree("pip-out/", ignore_errors=True)
102-
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
103-
for d in dirs:
104-
print(f"Cleaning {d}...")
105-
shutil.rmtree(d, ignore_errors=True)
106-
print("Done cleaning build artifacts.")
107-
sys.exit(0)
108-
elif arg == "--use-pt-pinned-commit":
118+
for pybind_arg in args.pybind:
119+
if pybind_arg not in ["coreml", "mps", "xnnpack"]:
120+
continue
121+
EXECUTORCH_BUILD_PYBIND = "ON"
122+
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
123+
124+
if args.clean:
125+
clean()
126+
return
127+
128+
if args.use_pt_pinned_commit:
109129
# This option is used in CI to make sure that PyTorch build from the pinned commit
110130
# is used instead of nightly. CI jobs wouldn't be able to catch regression from the
111131
# latest PT commit otherwise
112132
USE_PYTORCH_NIGHTLY = False
113-
else:
114-
print(f"Error: Unknown option {arg}")
115-
sys.exit(1)
116133

117-
# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
118-
# or is not turned off explicitly (--pybind off)
119-
# then install XNNPACK by default.
120-
if EXECUTORCH_BUILD_PYBIND == "":
121-
EXECUTORCH_BUILD_PYBIND = "ON"
122-
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"
123-
124-
# Use ClangCL on Windows.
125-
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
126-
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
127-
if os.name == "nt":
128-
CMAKE_ARGS += " -T ClangCL"
129-
130-
# Since ExecuTorch often uses main-branch features of pytorch, only the nightly
131-
# pip versions will have the required features.
132-
#
133-
# NOTE: If a newly-fetched version of the executorch repo changes the value of
134-
# NIGHTLY_VERSION, you should re-run this script to install the necessary
135-
# package versions.
136-
NIGHTLY_VERSION = "dev20250104"
137-
138-
# The pip repository that hosts nightly torch packages.
139-
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"
140-
141-
# pip packages needed by exir.
142-
EXIR_REQUIREMENTS = [
143-
# Setting USE_PYTORCH_NIGHTLY to false to test the pinned PyTorch commit. Note
144-
# that we don't need to set any version number there because they have already
145-
# been installed on CI before this step, so pip won't reinstall them
146-
f"torch==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torch",
147-
(
148-
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
149-
if USE_PYTORCH_NIGHTLY
150-
else "torchvision"
151-
), # For testing.
152-
"typing-extensions",
153-
]
154-
155-
# pip packages needed to run examples.
156-
# TODO: Make each example publish its own requirements.txt
157-
EXAMPLES_REQUIREMENTS = [
158-
"timm==1.0.7",
159-
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
160-
"torchsr==1.0.4",
161-
"transformers==4.47.1",
162-
]
163-
164-
# pip packages needed for development.
165-
DEVEL_REQUIREMENTS = [
166-
"cmake", # For building binary targets.
167-
"pip>=23", # For building the pip package.
168-
"pyyaml", # Imported by the kernel codegen tools.
169-
"setuptools>=63", # For building the pip package.
170-
"tomli", # Imported by extract_sources.py when using python < 3.11.
171-
"wheel", # For building the pip package archive.
172-
"zstd", # Imported by resolve_buck.py.
173-
"ai-edge-model-explorer>=0.1.16", # For visualizing ExportedPrograms
174-
]
175-
176-
# Assemble the list of requirements to actually install.
177-
# TODO: Add options for reducing the number of requirements.
178-
REQUIREMENTS_TO_INSTALL = EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS
179-
180-
# Install the requirements. `--extra-index-url` tells pip to look for package
181-
# versions on the provided URL if they aren't available on the default URL.
182-
subprocess.run(
183-
[
184-
sys.executable,
185-
"-m",
186-
"pip",
187-
"install",
188-
*REQUIREMENTS_TO_INSTALL,
189-
"--extra-index-url",
190-
TORCH_NIGHTLY_URL,
191-
],
192-
check=True,
193-
)
194-
195-
LOCAL_REQUIREMENTS = [
196-
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
197-
]
198-
199-
# Install packages directly from local copy instead of pypi.
200-
# This is usually not recommended.
201-
subprocess.run(
202-
[
203-
sys.executable,
204-
"-m",
205-
"pip",
206-
"install",
207-
*LOCAL_REQUIREMENTS,
208-
],
209-
check=True,
210-
)
134+
# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
135+
# or is not turned off explicitly (--pybind off)
136+
# then install XNNPACK by default.
137+
if EXECUTORCH_BUILD_PYBIND == "":
138+
EXECUTORCH_BUILD_PYBIND = "ON"
139+
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"
140+
141+
# Use ClangCL on Windows.
142+
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
143+
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
144+
if os.name == "nt":
145+
CMAKE_ARGS += " -T ClangCL"
146+
147+
# Since ExecuTorch often uses main-branch features of pytorch, only the nightly
148+
# pip versions will have the required features.
149+
#
150+
# NOTE: If a newly-fetched version of the executorch repo changes the value of
151+
# NIGHTLY_VERSION, you should re-run this script to install the necessary
152+
# package versions.
153+
NIGHTLY_VERSION = "dev20250104"
154+
155+
# The pip repository that hosts nightly torch packages.
156+
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"
157+
158+
# pip packages needed by exir.
159+
EXIR_REQUIREMENTS = [
160+
# Setting USE_PYTORCH_NIGHTLY to false to test the pinned PyTorch commit. Note
161+
# that we don't need to set any version number there because they have already
162+
# been installed on CI before this step, so pip won't reinstall them
163+
f"torch==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torch",
164+
(
165+
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
166+
if USE_PYTORCH_NIGHTLY
167+
else "torchvision"
168+
), # For testing.
169+
"typing-extensions",
170+
]
171+
172+
# pip packages needed to run examples.
173+
# TODO: Make each example publish its own requirements.txt
174+
EXAMPLES_REQUIREMENTS = [
175+
"timm==1.0.7",
176+
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
177+
"torchsr==1.0.4",
178+
"transformers==4.47.1",
179+
]
180+
181+
# pip packages needed for development.
182+
DEVEL_REQUIREMENTS = [
183+
"cmake", # For building binary targets.
184+
"pip>=23", # For building the pip package.
185+
"pyyaml", # Imported by the kernel codegen tools.
186+
"setuptools>=63", # For building the pip package.
187+
"tomli", # Imported by extract_sources.py when using python < 3.11.
188+
"wheel", # For building the pip package archive.
189+
"zstd", # Imported by resolve_buck.py.
190+
"ai-edge-model-explorer>=0.1.16", # For visualizing ExportedPrograms
191+
]
192+
193+
# Assemble the list of requirements to actually install.
194+
# TODO: Add options for reducing the number of requirements.
195+
REQUIREMENTS_TO_INSTALL = (
196+
EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS
197+
)
198+
199+
# Install the requirements. `--extra-index-url` tells pip to look for package
200+
# versions on the provided URL if they aren't available on the default URL.
201+
subprocess.run(
202+
[
203+
sys.executable,
204+
"-m",
205+
"pip",
206+
"install",
207+
*REQUIREMENTS_TO_INSTALL,
208+
"--extra-index-url",
209+
TORCH_NIGHTLY_URL,
210+
],
211+
check=True,
212+
)
213+
214+
LOCAL_REQUIREMENTS = [
215+
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
216+
]
217+
218+
# Install packages directly from local copy instead of pypi.
219+
# This is usually not recommended.
220+
subprocess.run(
221+
[
222+
sys.executable,
223+
"-m",
224+
"pip",
225+
"install",
226+
*LOCAL_REQUIREMENTS,
227+
],
228+
check=True,
229+
)
230+
231+
#
232+
# Install executorch pip package. This also makes `flatc` available on the path.
233+
# The --extra-index-url may be necessary if pyproject.toml has a dependency on a
234+
# pre-release or nightly version of a torch package.
235+
#
236+
237+
# Set environment variables
238+
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
239+
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
240+
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS
241+
242+
# Run the pip install command
243+
subprocess.run(
244+
[
245+
sys.executable,
246+
"-m",
247+
"pip",
248+
"install",
249+
".",
250+
"--no-build-isolation",
251+
"-v",
252+
"--extra-index-url",
253+
TORCH_NIGHTLY_URL,
254+
],
255+
check=True,
256+
)
211257

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

218-
# Set environment variables
219-
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
220-
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
221-
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS
222-
223-
# Run the pip install command
224-
subprocess.run(
225-
[
226-
sys.executable,
227-
"-m",
228-
"pip",
229-
"install",
230-
".",
231-
"--no-build-isolation",
232-
"-v",
233-
"--extra-index-url",
234-
TORCH_NIGHTLY_URL,
235-
],
236-
check=True,
237-
)
259+
if __name__ == "__main__":
260+
main(sys.argv[1:])

0 commit comments

Comments
 (0)