install_pytorch_wheels.py

#!/usr/bin/env python3
"""
install_pytorch_wheels.py

Installs PyTorch wheels from a pip index URL.

Usage (from repo root):
    python .github/scripts/install_pytorch_wheels.py --index-url <URL> --amdgpu-family <FAMILY> [OPTIONS]

Examples:
    # Install latest versions
    python .github/scripts/install_pytorch_wheels.py \
        --index-url <BASE_INDEX_URL>/whl \
        --amdgpu-family gfx1250

    # Install specific versions (matching ROCm builds)
    python .github/scripts/install_pytorch_wheels.py \
        --index-url <BASE_INDEX_URL>/whl \
        --amdgpu-family gfx1250 \
        --torch-version "2.10.0+devrocm7.12.0.dev0.849eec43b..." \
        --torchaudio-version "2.11.0a0+devrocm7.12.0.dev0.849eec43b..." \
        --torchvision-version "0.25.0a0+devrocm7.12.0.dev0.849eec43b..."
"""

import argparse
import re
import subprocess
import sys
import urllib.parse
import urllib.request


# Package configuration: (name, always_install)
PACKAGES = {
    "torch": True,
    "torchaudio": True,
    "torchvision": True,
    "triton": False,
    "rocm[devel]": True,
}
PYTORCH_PKGS = ["torch", "torchaudio", "torchvision", "triton"]


def print_banner(title: str) -> None:
    """Print a formatted banner."""
    print("=" * 50)
    print(title)
    print("=" * 50)


def build_package_spec(name: str, version: str | None) -> str:
    """Build a pip package spec (e.g., 'torch==2.10.0' or 'torch')."""
    return f"{name}=={version}" if version else name

def get_latest_package_version_for_rocm(
    index_url: str, package_name: str, rocm_version: str, required: bool = True,
    version_prefix: str | None = None,
) -> str | None:
    """Return latest package version containing rocm_version by parsing the index HTML.

    If version_prefix is set (e.g. "2.9"), only versions whose base part starts
    with that prefix are considered.
    """

    # Build the URL for this package's index page (e.g. .../gfx1250/torch/).
    rocm_tag = f"rocm{rocm_version}"
    url = f"{index_url.rstrip('/')}/{package_name}/"
    # Fetch the package index page; on failure (e.g. 404, timeout) fail if always_install, else return None.
    try:
        with urllib.request.urlopen(url, timeout=30) as resp:
            html = resp.read().decode("utf-8", errors="ignore")
    except Exception as e:
        print(f"Error: failed to fetch index for {package_name}: {e}", file=sys.stderr)
        sys.exit(1)
    # Parse wheel links: format is package-VERSION-...whl (e.g. torch-0.26.0a0+rocm7.12...-cp312-....whl).
    # Version can contain dots and + (URL-encoded as %2B), so we capture everything up to .whl.
    pattern = re.compile(
        re.escape(package_name) + r"-(.+?)\.whl",
        re.IGNORECASE,
    )
    all_suffixes = [m.group(1).strip() for m in pattern.finditer(html)]
    # Keep only wheels whose version string contains the requested ROCm tag (e.g. rocm7.12.0a20260224).
    # Version is the first segment before "-" in the suffix; decode %2B to + for comparison.
    matching = []
    for s in all_suffixes:
        ver = s.split("-")[0]
        if rocm_tag in ver:
            matching.append(urllib.parse.unquote(ver))
    # Filter by version prefix (e.g. "2.9" matches "2.9.0+...", "2.9.1+...").
    if version_prefix and matching:
        matching = [v for v in matching if v.split("+")[0].startswith(version_prefix)]
    # No matching wheels: if required (always_install), fail; otherwise return None (package will be skipped).
    if not matching:
        if required:
            msg = f"Error: no wheel found for {package_name} with ROCm {rocm_version}"
            if version_prefix:
                msg += f" and version prefix {version_prefix}"
            print(msg, file=sys.stderr)
            sys.exit(1)
        return None
    # Pick the latest version by comparing all numeric parts including the ROCm date.
    def _key(v: str) -> tuple[int, ...]:
        try:
            return tuple(int(x) for x in re.split(r"[.\-a+]", v) if x.isdigit())
        except (ValueError, AttributeError):
            return (0,)
    return max(matching, key=_key)


def run_pip_install(
    index_url: str, packages: list[str], break_system_packages: bool = True
) -> None:
    """Run pip install with the given packages."""
    cmd = [sys.executable, "-m", "pip", "install", "--index-url", index_url]

    if break_system_packages:
        cmd.append("--break-system-packages")

    cmd.extend(packages)

    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, check=False)

    if result.returncode != 0:
        print(f"Error: pip install failed with return code {result.returncode}")
        sys.exit(result.returncode)


def check_package(name: str) -> tuple[bool, str | None]:
    """Check if a package is installed and return (installed, version)."""
    try:
        module = __import__(name)
        return True, getattr(module, "__version__", "unknown")
    except ImportError:
        return False, None


def verify_installation() -> bool:
    """Verify PyTorch installation and print version info."""
    print_banner("Verifying Installation")

    # Check torch separately for ROCm info
    try:
        import torch as _torch

        version = getattr(_torch, "__version__", "unknown")
    except ImportError as e:
        print(f"Error: torch import failed ({e!r}). If wheels are installed, run rocm-sdk init first.")
        return False

    print(f"torch: {version}")

    hip_version = _torch.version.hip
    print(f"ROCm/HIP: {hip_version or 'not available'}")
    print(f"Built with ROCm: {hip_version is not None}")

    # Check other packages
    for name in ["torchaudio", "torchvision", "triton", "rocm"]:
        installed, version = check_package(name)
        status = version if installed else "not installed"
        print(f"{name}: {status}")

    return True


def list_installed_packages() -> None:
    """List installed torch-related packages."""
    print("\nInstalled PyTorch packages:")
    result = subprocess.run(
        [sys.executable, "-m", "pip", "list"],
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode == 0:
        keywords = ["torch", "triton", "rocm"]
        for line in result.stdout.splitlines():
            if any(kw in line.lower() for kw in keywords):
                print(f"  {line}")


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Install PyTorch wheels from a pip index URL",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        "--index-url", required=True, help="Base URL for PyTorch wheels index"
    )
    parser.add_argument(
        "--amdgpu-family", required=True, help="AMD GPU family (e.g., gfx1250)"
    )
    parser.add_argument(
        "--rocm-version",
        help="Optional. ROCm version (e.g. 7.12.0a20260126). When set without --torch-version: discovers and installs latest torch/torchaudio/torchvision/triton built for this ROCm. ",
    )
    parser.add_argument(
        "--torch-version", help="Specific torch version (default: latest)"
    )
    parser.add_argument(
        "--torch-version-prefix",
        help="Torch version prefix for discovery (e.g. '2.9' matches 2.9.x). "
             "Only used in auto-discovery mode (--rocm-version without --torch-version).",
    )
    parser.add_argument(
        "--torchaudio-version", help="Specific torchaudio version (default: latest)"
    )
    parser.add_argument(
        "--torchvision-version", help="Specific torchvision version (default: latest)"
    )
    parser.add_argument(
        "--triton-version",
        help="Specific triton version (default: from torch dependency)",
    )
    parser.add_argument(
        "--no-break-system-packages",
        action="store_true",
        help="Don't use --break-system-packages",
    )
    parser.add_argument(
        "--skip-verify", action="store_true", help="Skip verification step"
    )

    args = parser.parse_args()

    # Build the full index URL
    index_url = f"{args.index_url.rstrip('/')}/{args.amdgpu_family}/"

    rocm = args.rocm_version
    rocm_only = bool(rocm and not args.torch_version)
    torch_prefix = args.torch_version_prefix if rocm_only else None
    break_sys = not args.no_break_system_packages

    if rocm_only:
        # Two-pass install:
        #   Pass 1: torch (pinned) + rocm[devel] (pinned)
        #   Pass 2: torchaudio, torchvision, triton (unpinned — pip resolves compatibility)
        torch_version = get_latest_package_version_for_rocm(
            index_url, "torch", rocm, required=True, version_prefix=torch_prefix,
        )

        print_banner("PyTorch Wheels Installation")
        print(f"Index URL:      {index_url}")
        print(f"AMDGPU Family:  {args.amdgpu_family}")
        print(f"Python:         {sys.version_info.major}.{sys.version_info.minor}")
        print(f"torch:          {torch_version}")
        print(f"rocm[devel]:    {rocm}")
        print(f"torchaudio:     (pip resolves)")
        print(f"torchvision:    (pip resolves)")
        print(f"triton:         (torch dependency)")
        print("=" * 50)

        # Pass 1: install torch + rocm[devel] with exact versions.
        # torch's declared dependency on triton pulls in the correct build.
        primary = [
            build_package_spec("torch", torch_version),
            build_package_spec("rocm[devel]", rocm),
        ]
        print_banner("Pass 1: torch + rocm[devel]")
        print(f"Installing: {', '.join(primary)}")
        run_pip_install(index_url, primary, break_sys)

        # Pass 2: install torchaudio/torchvision without pinning — pip picks
        # versions compatible with the torch that's already installed
        companions = ["torchaudio", "torchvision"]
        print_banner("Pass 2: torchaudio, torchvision (unpinned)")
        print(f"Installing: {', '.join(companions)}")
        run_pip_install(index_url, companions, break_sys)
    else:
        # Explicit versions mode — install everything in one shot
        arg_attrs = ["torch_version", "torchaudio_version", "torchvision_version", "triton_version"]
        versions = {p: getattr(args, a) for p, a in zip(PYTORCH_PKGS, arg_attrs)}
        versions["rocm[devel]"] = rocm if rocm else None

        print_banner("PyTorch Wheels Installation")
        print(f"Index URL:      {index_url}")
        print(f"AMDGPU Family:  {args.amdgpu_family}")
        print(f"Python:         {sys.version_info.major}.{sys.version_info.minor}")
        for name, version in versions.items():
            print(f"{name:14}: {version or 'latest'}")
        print("=" * 50)

        packages = []
        for name, always_install in PACKAGES.items():
            version = versions.get(name)
            if always_install or version:
                packages.append(build_package_spec(name, version))

        print(f"Installing: {', '.join(packages)}")
        run_pip_install(index_url, packages, break_sys)

    # Verify
    if not args.skip_verify and not verify_installation():
        return 1

    list_installed_packages()
    print_banner("Installation complete")
    return 0


if __name__ == "__main__":
    sys.exit(main())