install_requirements.py: reorganize requirements

install_requirements.py

@@ -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
@@ -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.
@@ -135,103 +88,167 @@ 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.
+    ]
+
+    EXAMPLES_REQUIREMENTS = [
+        f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if use_pytorch_nightly else "torchaudio",
+    ]
+
+    # Assemble the list of requirements to actually install.
+    # TODO: Add options for reducing the number of requirements.
+    REQUIREMENTS_TO_INSTALL = (
+        EXIR_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",
+            "-r",
+            "requirements-dev.txt",
+            "-r",
+            "requirements-examples.txt",
+            *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",
+            # Without --no-build-isolation, setup.py can't find the torch module.
+            "--no-build-isolation",
+            *LOCAL_REQUIREMENTS,
+        ],
+        check=True,
+    )
+
+
+def main(args):
+    if not python_is_compatible():
+        sys.exit(1)
 
-#
-# 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.
-#
+    # 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
 
-# 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 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,
+    )
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

requirements-dev.txt

@@ -0,0 +1,2 @@
+# For visualizing ExportedPrograms
+ai-edge-model-explorer >= 0.1.16

requirements-examples.txt

@@ -0,0 +1,5 @@
+# pip packages needed to run examples.
+# TODO: Make each example publish its own requirements.txt
+timm == 1.0.7
+torchsr == 1.0.4
+transformers ==4.47.1