Skip to content

Commit 007ea3e

Browse files
swolchokswolchok
authored
install_requirements.py: use argparse, minor cleanup (#7703)
Replace the bespoke argument parser with argparse in preparation for separating requirements installation from building and installing ExecuTorch itself.
1 parent 2f0518d commit 007ea3e

File tree

1 file changed

+188
-159
lines changed

1 file changed

+188
-159
lines changed

‎install_requirements.py

+188-159
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,174 +65,201 @@ 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+
VALID_PYBINDS = ["coreml", "mps", "xnnpack"]
7580

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)
81+
82+
def main(args):
83+
if not python_is_compatible():
84+
sys.exit(1)
85+
86+
# Parse options.
87+
88+
EXECUTORCH_BUILD_PYBIND = ""
89+
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
90+
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
91+
USE_PYTORCH_NIGHTLY = True
92+
93+
parser = argparse.ArgumentParser()
94+
parser.add_argument(
95+
"--pybind",
96+
action="append",
97+
nargs="+",
98+
help="one or more of coreml/mps/xnnpack, or off",
99+
)
100+
parser.add_argument(
101+
"--clean",
102+
action="store_true",
103+
help="clean build artifacts and pip-out instead of installing",
104+
)
105+
parser.add_argument(
106+
"--use-pt-pinned-commit",
107+
action="store_true",
108+
help="build from the pinned PyTorch commit instead of nightly",
109+
)
110+
args = parser.parse_args(args)
111+
112+
if args.clean:
113+
clean()
114+
return
115+
116+
if args.pybind:
117+
# Flatten list of lists.
118+
args.pybind = list(itertools.chain(*args.pybind))
119+
if "off" in args.pybind:
120+
if len(args.pybind) != 1:
121+
raise Exception(
122+
f"Cannot combine `off` with other pybinds: {args.pybind}"
123+
)
93124
EXECUTORCH_BUILD_PYBIND = "OFF"
94125
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":
126+
for pybind_arg in args.pybind:
127+
if pybind_arg not in VALID_PYBINDS:
128+
raise Exception(
129+
f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
130+
)
131+
EXECUTORCH_BUILD_PYBIND = "ON"
132+
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
133+
134+
if args.use_pt_pinned_commit:
109135
# This option is used in CI to make sure that PyTorch build from the pinned commit
110136
# is used instead of nightly. CI jobs wouldn't be able to catch regression from the
111137
# latest PT commit otherwise
112138
USE_PYTORCH_NIGHTLY = False
113-
else:
114-
print(f"Error: Unknown option {arg}")
115-
sys.exit(1)
116139

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

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

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

0 commit comments

Comments
 (0)