install_requirements.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# Copyright 2024-25 Arm Limited and/or its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import os
import platform
import subprocess
import sys

from install_utils import determine_torch_url, is_intel_mac_os, python_is_compatible

# The pip repository that hosts nightly torch packages.
# This will be dynamically set based on CUDA availability and CUDA backend enabled/disabled.
TORCH_URL_BASE = "https://download.pytorch.org/whl/test"
TORCHAO_URL_BASE = "https://download.pytorch.org/whl/nightly"
TORCHAO_NIGHTLY_VERSION = "0.18.0.dev20260729"
CU134_TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260811"
# These wheels' metadata pairs August 11 domain libraries with August 10 torch.
CU134_TORCH_PACKAGES = [
    "torch==2.14.0.dev20260810+cu134",
    "torchvision==0.29.0.dev20260811+cu134",
    "torchaudio==2.11.0.dev20260811+cu134",
]


def torchao_from_source():
    return (
        os.environ.get("EXECUTORCH_BUILD_KERNELS_TORCHAO") == "1"
        or os.environ.get("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1"
    )


def cu134_requirements(torch_url, include_domains=False):
    if not torch_url.endswith("/cu134"):
        return []
    packages = list(
        CU134_TORCH_PACKAGES if include_domains else CU134_TORCH_PACKAGES[:1]
    )
    if not torchao_from_source():
        torchao_variant = (
            "cpu" if platform.machine().lower() in ("aarch64", "arm64") else "cu134"
        )
        packages.append(f"torchao=={CU134_TORCHAO_NIGHTLY_VERSION}+{torchao_variant}")
    return packages


# 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.
#
# NOTE: If you're changing, make the corresponding change in .ci/docker/ci_commit_pins/pytorch.txt
# by picking the hash from the same date in
# https://hud.pytorch.org/hud/pytorch/pytorch/nightly/ @lint-ignore
#
# NOTE: If you're changing, make the corresponding supported CUDA versions in
# SUPPORTED_CUDA_VERSIONS in install_utils.py if needed.


def install_requirements(use_pytorch_nightly):
    # Skip pip install on Intel macOS if using nightly.
    if use_pytorch_nightly and is_intel_mac_os():
        print(
            "ERROR: Prebuilt PyTorch wheels are no longer available for Intel-based macOS.\n"
            "Please build from source by following https://docs.pytorch.org/executorch/main/using-executorch-building-from-source.html",
            file=sys.stderr,
        )
        sys.exit(1)

    # Determine the appropriate PyTorch URL based on CUDA delegate status
    torch_url = determine_torch_url(TORCH_URL_BASE)
    cu134_packages = cu134_requirements(torch_url)
    if cu134_packages:
        torch_url = determine_torch_url(TORCHAO_URL_BASE)
        if not use_pytorch_nightly:
            cu134_packages[0] = "torch"
    # torchao's CUDA channel publishes x86_64 only, so asking for a CUDA build makes the pin
    # unsatisfiable on aarch64. Only that case is special-cased: falling back everywhere would
    # change which torchao a CPU x86_64 install resolves, and the CUDA build is genuinely wanted
    # where it exists. This nightly is what a development checkout is tested against, and the
    # wheel's own torchao lower bound is this same version so that installing the package
    # afterwards leaves this pin in place rather than replacing it.
    if platform.machine().lower() in ("aarch64", "arm64"):
        # The cpu channel specifically, not the index root. The root carries every variant, and a
        # pin without a local segment admits all of them while ordering a local segment highest,
        # so the xpu channel's pure python wheel would win on version before pip compares wheel
        # tags, silently replacing the compiled aarch64 build.
        torchao_url = f"{TORCHAO_URL_BASE}/cpu"
    else:
        torchao_url = determine_torch_url(TORCHAO_URL_BASE)

    # pip packages needed by exir.
    TORCH_PACKAGE = cu134_packages or [
        # 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
        ("torch==2.14.0" if use_pytorch_nightly else "torch"),
        f"torchao=={TORCHAO_NIGHTLY_VERSION}",
    ]

    # Install the requirements for core ExecuTorch package.
    # `--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",
            *TORCH_PACKAGE,
            "--extra-index-url",
            torch_url,
            "--extra-index-url",
            torchao_url,
        ],
        check=True,
    )

    LOCAL_REQUIREMENTS = []
    if torchao_from_source():
        LOCAL_REQUIREMENTS.append("third-party/ao")
    if sys.platform != "win32":
        # TODO(larryliu0820): Setup a pypi package for this.
        LOCAL_REQUIREMENTS.append("extension/llm/tokenizers")
    # TODO(gjcomer): Re-enable when buildable on Windows.

    if LOCAL_REQUIREMENTS:
        # Install packages directly from local copy instead of pypi.
        # This is usually not recommended.
        new_env = os.environ.copy()
        if ("EXECUTORCH_BUILD_KERNELS_TORCHAO" not in new_env) or (
            new_env["EXECUTORCH_BUILD_KERNELS_TORCHAO"] == "0"
        ):
            new_env["USE_CPP"] = "0"
        else:
            assert new_env["EXECUTORCH_BUILD_KERNELS_TORCHAO"] == "1"
            new_env["USE_CPP"] = "1"
            new_env["CMAKE_POLICY_VERSION_MINIMUM"] = "3.5"
        subprocess.run(
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                # Without --no-build-isolation, setup.py can't find the torch module.
                "--no-build-isolation",
                *LOCAL_REQUIREMENTS,
                *cu134_packages,
                *(
                    ["--extra-index-url", torch_url, "--extra-index-url", torchao_url]
                    if cu134_packages
                    else []
                ),
            ],
            env=new_env,
            check=True,
        )


def install_optional_example_requirements(use_pytorch_nightly):
    # Determine the appropriate PyTorch URL based on CUDA delegate status
    torch_url = determine_torch_url(TORCH_URL_BASE)
    cu134_packages = (
        cu134_requirements(torch_url, include_domains=True)
        if use_pytorch_nightly
        else []
    )
    if cu134_packages:
        torch_url = determine_torch_url(TORCHAO_URL_BASE)
    torchao_index = (
        ["--extra-index-url", f"{TORCHAO_URL_BASE}/cpu"]
        if cu134_packages and platform.machine().lower() in ("aarch64", "arm64")
        else []
    )

    print("Installing torch domain libraries")
    DOMAIN_LIBRARIES = cu134_packages or [
        ("torchvision==0.29.0" if use_pytorch_nightly else "torchvision"),
        ("torchaudio==2.11.0" if use_pytorch_nightly else "torchaudio"),
    ]
    # Then install domain libraries
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            *DOMAIN_LIBRARIES,
            "--extra-index-url",
            torch_url,
            *torchao_index,
        ],
        check=True,
    )

    print("Installing packages in requirements-examples.txt")
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-r",
            "requirements-examples.txt",
            *cu134_packages,
            "--extra-index-url",
            torch_url,
            *torchao_index,
            "--upgrade-strategy",
            "only-if-needed",
        ],
        check=True,
    )


def main(args):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--use-pt-pinned-commit",
        action="store_true",
        help="build from the pinned PyTorch commit instead of nightly",
    )
    parser.add_argument(
        "--example",
        action="store_true",
        help="Also installs required packages for running example scripts.",
    )
    args = parser.parse_args(args)
    use_pytorch_nightly = not bool(args.use_pt_pinned_commit)
    install_requirements(use_pytorch_nightly)
    if args.example:
        install_optional_example_requirements(use_pytorch_nightly)


if __name__ == "__main__":
    # Before doing anything, cd to the directory containing this script.
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    if not python_is_compatible():
        sys.exit(1)
    main(sys.argv[1:])