install_utils.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 functools
import platform
import re
import subprocess
import sys
from typing import List, Optional

# Supported CUDA versions - modify this to add/remove supported versions
# Format: tuple of (major, minor) version numbers
SUPPORTED_CUDA_VERSIONS = (
    (12, 6),
    (12, 8),
    (12, 9),
    (13, 0),
)


def is_cmake_option_on(
    cmake_configuration_args: List[str], var_name: str, default: bool
) -> bool:
    """
    Get a boolean CMake variable, from a list of CMake configuration arguments.
    The var_name should not include the "-D" prefix.

    Args:
        cmake_configuration_args: List of CMake configuration arguments.
        var_name: Name of the CMake variable.
        default: Default boolean value if the variable is not set.

    Returns:
        Boolean value of the CMake variable.
    """
    cmake_define = _extract_cmake_define(cmake_configuration_args, var_name)

    return _normalize_cmake_bool(cmake_define, default)


def is_cuda_available() -> bool:
    """
    Check if CUDA is available on the system by attempting to get the CUDA version.

    Returns:
        True if CUDA is available and supported, False otherwise.
    """
    try:
        _get_cuda_version()
        return True
    except Exception:
        return False


@functools.lru_cache(maxsize=1)
def _get_cuda_version():
    """
    Get the CUDA version installed on the system using nvcc command.
    Returns a tuple (major, minor).

    Raises:
        RuntimeError: if nvcc is not found or version cannot be parsed
    """
    try:
        # Get CUDA version from nvcc (CUDA compiler)
        nvcc_result = subprocess.run(
            ["nvcc", "--version"], capture_output=True, text=True, check=True
        )
        # Parse nvcc output for CUDA version
        # Output contains line like "Cuda compilation tools, release 12.6, V12.6.68"
        match = re.search(r"release (\d+)\.(\d+)", nvcc_result.stdout)
        if match:
            major, minor = int(match.group(1)), int(match.group(2))

            # Check if the detected version is supported
            if (major, minor) not in SUPPORTED_CUDA_VERSIONS:
                available_versions = ", ".join(
                    [f"{maj}.{min}" for maj, min in SUPPORTED_CUDA_VERSIONS]
                )
                raise RuntimeError(
                    f"Detected CUDA version {major}.{minor} is not supported. "
                    f"Supported versions: {available_versions}."
                )

            return (major, minor)
        else:
            raise RuntimeError(
                "Failed to parse CUDA version from nvcc output. "
                "Ensure CUDA is properly installed."
            )
    except FileNotFoundError:
        raise RuntimeError(
            "nvcc (CUDA compiler) is not found in PATH. Install the CUDA toolkit."
        )
    except subprocess.CalledProcessError as e:
        raise RuntimeError(
            f"nvcc command failed with error: {e}. "
            "Ensure CUDA is properly installed."
        )


def _extract_cmake_define(args: List[str], name: str) -> Optional[str]:
    prefix = f"-D{name}="
    for arg in args:
        if arg.startswith(prefix):
            return arg[len(prefix) :]
    return None


def _normalize_cmake_bool(value: Optional[str], default: bool = False) -> bool:
    if value is None:
        return default
    normalized = value.strip().upper()
    if normalized in {"ON", "1", "TRUE", "YES"}:
        return True
    if normalized in {"OFF", "0", "FALSE", "NO"}:
        return False
    return default


def _cuda_version_to_pytorch_suffix(major, minor):
    """
    Generate PyTorch CUDA wheel suffix from CUDA version numbers.

    Args:
        major: CUDA major version (e.g., 12)
        minor: CUDA minor version (e.g., 6)

    Returns:
        PyTorch wheel suffix string (e.g., "cu126")
    """
    return f"cu{major}{minor}"


def _get_pytorch_cuda_url(cuda_version, torch_nightly_url_base):
    """
    Get the appropriate PyTorch CUDA URL for the given CUDA version.

    Args:
        cuda_version: tuple of (major, minor) version numbers
        torch_nightly_url_base: Base URL for PyTorch nightly packages

    Returns:
        URL string for PyTorch CUDA packages
    """
    major, minor = cuda_version
    # Generate CUDA suffix (version validation is already done in _get_cuda_version)
    cuda_suffix = _cuda_version_to_pytorch_suffix(major, minor)

    return f"{torch_nightly_url_base}/{cuda_suffix}"


@functools.lru_cache(maxsize=1)
def determine_torch_url(torch_nightly_url_base):
    """
    Determine the appropriate PyTorch installation URL based on CUDA availability.
    Uses @functools.lru_cache to avoid redundant CUDA detection and print statements.

    Args:
        torch_nightly_url_base: Base URL for PyTorch nightly packages

    Returns:
        URL string for PyTorch packages
    """
    if platform.system().lower() == "windows":
        print(
            "Windows detected, using CPU-only PyTorch until CUDA support is available"
        )
        return f"{torch_nightly_url_base}/cpu"

    print("Attempting to detect CUDA via nvcc...")

    try:
        cuda_version = _get_cuda_version()
    except Exception as err:
        print(f"CUDA detection failed ({err}), using CPU-only PyTorch")
        return f"{torch_nightly_url_base}/cpu"

    major, minor = cuda_version
    print(f"Detected CUDA version: {major}.{minor}")

    # Get appropriate PyTorch CUDA URL
    torch_url = _get_pytorch_cuda_url(cuda_version, torch_nightly_url_base)
    print(f"Using PyTorch URL: {torch_url}")

    return torch_url


# Prebuilt binaries for Intel-based macOS are no longer available on PyPI; users must compile from source.
# PyTorch stopped building macOS x86_64 binaries since version 2.3.0 (January 2024).
def is_intel_mac_os():
    # Returns True if running on Intel macOS.
    return platform.system().lower() == "darwin" and platform.machine().lower() in (
        "x86",
        "x86_64",
        "i386",
    )


def python_is_compatible():
    # Scrape the version range from pyproject.toml, which should be in the current directory.
    version_specifier = None
    with open("pyproject.toml", "r") as file:
        for line in file:
            if line.startswith("requires-python"):
                match = re.search(r'"([^"]*)"', line)
                if match:
                    version_specifier = match.group(1)
                    break

    if not version_specifier:
        print(
            "WARNING: Skipping python version check: version range not found",
            file=sys.stderr,
        )
        return False

    # Install the packaging module if necessary.
    try:
        import packaging
    except ImportError:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "packaging"], check=True
        )
    # Compare the current python version to the range in version_specifier. Exits
    # with status 1 if the version is not compatible, or with status 0 if the
    # version is compatible or the logic itself fails.
    try:
        import packaging.specifiers
        import packaging.version

        python_version = packaging.version.parse(platform.python_version())
        version_range = packaging.specifiers.SpecifierSet(version_specifier)
        if python_version not in version_range:
            print(
                f'ERROR: ExecuTorch does not support python version {python_version}: must satisfy "{version_specifier}"',
                file=sys.stderr,
            )
            return False
    except Exception as e:
        print(f"WARNING: Skipping python version check: {e}", file=sys.stderr)
    return True