.ci/aarch64_linux/aarch64_ci_build.sh

@@ -3,22 +3,15 @@ set -eux -o pipefail
 
 GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-}
 
+if [[ "$GPU_ARCH_VERSION" == *"12.6"* ]]; then
+    export TORCH_CUDA_ARCH_LIST="9.0"
+elif [[ "$GPU_ARCH_VERSION" == *"12.8"* ]]; then
+    export TORCH_CUDA_ARCH_LIST="9.0;10.0;12.0"
+fi
+
 SCRIPTPATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 source $SCRIPTPATH/aarch64_ci_setup.sh
 
-tagged_version() {
-  GIT_DESCRIBE="git --git-dir /pytorch/.git describe --tags --match v[0-9]*.[0-9]*.[0-9]*"
-  if ${GIT_DESCRIBE} --exact >/dev/null; then
-    ${GIT_DESCRIBE}
-  else
-    return 1
-  fi
-}
-
-if tagged_version >/dev/null; then
-  export OVERRIDE_PACKAGE_VERSION="$(tagged_version | sed -e 's/^v//' -e 's/-.*$//')"
-fi
-
 ###############################################################################
 # Run aarch64 builder python
 ###############################################################################
@@ -27,7 +20,7 @@ cd /
 # on the mounted pytorch repo
 git config --global --add safe.directory /pytorch
 pip install -r /pytorch/requirements.txt
-pip install auditwheel
+pip install auditwheel==6.2.0
 if [ "$DESIRED_CUDA" = "cpu" ]; then
     echo "BASE_CUDA_VERSION is not set. Building cpu wheel."
     #USE_PRIORITIZED_TEXT_FOR_LD for enable linker script optimization https://github.com/pytorch/pytorch/pull/121975/files

.ci/aarch64_linux/aarch64_ci_setup.sh

@@ -5,16 +5,14 @@ set -eux -o pipefail
 # By creating symlinks from desired /opt/python to /usr/local/bin/
 
 NUMPY_VERSION=2.0.2
-PYGIT2_VERSION=1.15.1
-if [[ "$DESIRED_PYTHON"  == "3.13" ]]; then
+if [[ "$DESIRED_PYTHON"  == "3.13" || "$DESIRED_PYTHON" == "3.13t" ]]; then
     NUMPY_VERSION=2.1.2
-    PYGIT2_VERSION=1.16.0
 fi
 
 SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
 source $SCRIPTPATH/../manywheel/set_desired_python.sh
 
-pip install -q numpy==${NUMPY_VERSION} pyyaml==6.0.2 scons==4.7.0 ninja==1.11.1 patchelf==0.17.2 pygit2==${PYGIT2_VERSION}
+pip install -q numpy==${NUMPY_VERSION} pyyaml==6.0.2 scons==4.7.0 ninja==1.11.1 patchelf==0.17.2
 
 for tool in python python3 pip pip3 ninja scons patchelf; do
     ln -sf ${DESIRED_PYTHON_BIN_DIR}/${tool} /usr/local/bin;

.ci/aarch64_linux/aarch64_wheel_ci_build.py

@@ -4,12 +4,9 @@
 import os
 import shutil
 from subprocess import check_call, check_output
-from typing import List
 
-from pygit2 import Repository
 
-
-def list_dir(path: str) -> List[str]:
+def list_dir(path: str) -> list[str]:
     """'
     Helper for getting paths for Python
     """
@@ -42,7 +39,7 @@ def build_ArmComputeLibrary() -> None:
             "clone",
             "https://github.com/ARM-software/ComputeLibrary.git",
             "-b",
-            "v24.09",
+            "v25.02",
             "--depth",
             "1",
             "--shallow-submodules",
@@ -58,7 +55,7 @@ def build_ArmComputeLibrary() -> None:
         shutil.copytree(f"{acl_checkout_dir}/{d}", f"{acl_install_dir}/{d}")
 
 
-def update_wheel(wheel_path) -> None:
+def update_wheel(wheel_path, desired_cuda) -> None:
     """
     Update the cuda wheel libraries
     """
@@ -80,7 +77,6 @@ def update_wheel(wheel_path) -> None:
         "/usr/local/cuda/lib64/libnvToolsExt.so.1",
         "/usr/local/cuda/lib64/libnvJitLink.so.12",
         "/usr/local/cuda/lib64/libnvrtc.so.12",
-        "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.6",
         "/usr/local/cuda/lib64/libcudnn_adv.so.9",
         "/usr/local/cuda/lib64/libcudnn_cnn.so.9",
         "/usr/local/cuda/lib64/libcudnn_graph.so.9",
@@ -100,6 +96,18 @@ def update_wheel(wheel_path) -> None:
             "/usr/local/lib/libnvpl_lapack_core.so.0",
             "/usr/local/lib/libnvpl_blas_core.so.0",
         ]
+        if "126" in desired_cuda:
+            libs_to_copy += [
+                "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.6",
+                "/usr/local/cuda/lib64/libcufile.so.0",
+                "/usr/local/cuda/lib64/libcufile_rdma.so.1",
+            ]
+        elif "128" in desired_cuda:
+            libs_to_copy += [
+                "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.8",
+                "/usr/local/cuda/lib64/libcufile.so.0",
+                "/usr/local/cuda/lib64/libcufile_rdma.so.1",
+            ]
     else:
         libs_to_copy += [
             "/opt/OpenBLAS/lib/libopenblas.so.0",
@@ -171,22 +179,22 @@ def parse_arguments():
     args = parse_arguments()
     enable_mkldnn = args.enable_mkldnn
     enable_cuda = args.enable_cuda
-    repo = Repository("/pytorch")
-    branch = repo.head.name
-    if branch == "HEAD":
-        branch = "master"
+    branch = check_output(
+        ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd="/pytorch"
+    ).decode()
 
     print("Building PyTorch wheel")
     build_vars = "MAX_JOBS=5 CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
     os.system("cd /pytorch; python setup.py clean")
 
     override_package_version = os.getenv("OVERRIDE_PACKAGE_VERSION")
+    desired_cuda = os.getenv("DESIRED_CUDA")
     if override_package_version is not None:
         version = override_package_version
         build_vars += (
             f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version} PYTORCH_BUILD_NUMBER=1 "
         )
-    elif branch in ["nightly", "master"]:
+    elif branch in ["nightly", "main"]:
         build_date = (
             check_output(["git", "log", "--pretty=format:%cs", "-1"], cwd="/pytorch")
             .decode()
@@ -196,12 +204,11 @@ def parse_arguments():
             check_output(["cat", "version.txt"], cwd="/pytorch").decode().strip()[:-2]
         )
         if enable_cuda:
-            desired_cuda = os.getenv("DESIRED_CUDA")
             build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date}+{desired_cuda} PYTORCH_BUILD_NUMBER=1 "
         else:
             build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1 "
     elif branch.startswith(("v1.", "v2.")):
-        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1 "
+        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1 : branch.find('-')]} PYTORCH_BUILD_NUMBER=1 "
 
     if enable_mkldnn:
         build_ArmComputeLibrary()
@@ -225,6 +232,6 @@ def parse_arguments():
         print("Updating Cuda Dependency")
         filename = os.listdir("/pytorch/dist/")
         wheel_path = f"/pytorch/dist/{filename[0]}"
-        update_wheel(wheel_path)
+        update_wheel(wheel_path, desired_cuda)
     pytorch_wheel_name = complete_wheel("/pytorch/")
     print(f"Build Complete. Created {pytorch_wheel_name}..")

.ci/aarch64_linux/build_aarch64_wheel.py

@@ -12,22 +12,22 @@
 import subprocess
 import sys
 import time
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import boto3
 
 
 # AMI images for us-east-1, change the following based on your ~/.aws/config
 os_amis = {
-    "ubuntu18_04": "ami-078eece1d8119409f",  # login_name: ubuntu
     "ubuntu20_04": "ami-052eac90edaa9d08f",  # login_name: ubuntu
     "ubuntu22_04": "ami-0c6c29c5125214c77",  # login_name: ubuntu
     "redhat8": "ami-0698b90665a2ddcf1",  # login_name: ec2-user
 }
-ubuntu18_04_ami = os_amis["ubuntu18_04"]
 
+ubuntu20_04_ami = os_amis["ubuntu20_04"]
 
-def compute_keyfile_path(key_name: Optional[str] = None) -> Tuple[str, str]:
+
+def compute_keyfile_path(key_name: Optional[str] = None) -> tuple[str, str]:
     if key_name is None:
         key_name = os.getenv("AWS_KEY_NAME")
         if key_name is None:
@@ -57,7 +57,7 @@ def ec2_instances_by_id(instance_id):
 
 
 def start_instance(
-    key_name, ami=ubuntu18_04_ami, instance_type="t4g.2xlarge", ebs_size: int = 50
+    key_name, ami=ubuntu20_04_ami, instance_type="t4g.2xlarge", ebs_size: int = 50
 ):
     inst = ec2.create_instances(
         ImageId=ami,
@@ -96,7 +96,7 @@ def __init__(self, addr: str, keyfile_path: str, login_name: str = "ubuntu"):
         self.keyfile_path = keyfile_path
         self.login_name = login_name
 
-    def _gen_ssh_prefix(self) -> List[str]:
+    def _gen_ssh_prefix(self) -> list[str]:
         return [
             "ssh",
             "-o",
@@ -108,13 +108,13 @@ def _gen_ssh_prefix(self) -> List[str]:
         ]
 
     @staticmethod
-    def _split_cmd(args: Union[str, List[str]]) -> List[str]:
+    def _split_cmd(args: Union[str, list[str]]) -> list[str]:
         return args.split() if isinstance(args, str) else args
 
-    def run_ssh_cmd(self, args: Union[str, List[str]]) -> None:
+    def run_ssh_cmd(self, args: Union[str, list[str]]) -> None:
         subprocess.check_call(self._gen_ssh_prefix() + self._split_cmd(args))
 
-    def check_ssh_output(self, args: Union[str, List[str]]) -> str:
+    def check_ssh_output(self, args: Union[str, list[str]]) -> str:
         return subprocess.check_output(
             self._gen_ssh_prefix() + self._split_cmd(args)
         ).decode("utf-8")
@@ -157,7 +157,7 @@ def start_docker(self, image="quay.io/pypa/manylinux2014_aarch64:latest") -> Non
     def using_docker(self) -> bool:
         return self.container_id is not None
 
-    def run_cmd(self, args: Union[str, List[str]]) -> None:
+    def run_cmd(self, args: Union[str, list[str]]) -> None:
         if not self.using_docker():
             return self.run_ssh_cmd(args)
         assert self.container_id is not None
@@ -178,7 +178,7 @@ def run_cmd(self, args: Union[str, List[str]]) -> None:
         if rc != 0:
             raise subprocess.CalledProcessError(rc, docker_cmd)
 
-    def check_output(self, args: Union[str, List[str]]) -> str:
+    def check_output(self, args: Union[str, list[str]]) -> str:
         if not self.using_docker():
             return self.check_ssh_output(args)
         assert self.container_id is not None
@@ -230,7 +230,7 @@ def download_wheel(
             )
         self.download_file(remote_file, local_file)
 
-    def list_dir(self, path: str) -> List[str]:
+    def list_dir(self, path: str) -> list[str]:
         return self.check_output(["ls", "-1", path]).split("\n")
 
 
@@ -327,7 +327,7 @@ def build_ArmComputeLibrary(host: RemoteHost, git_clone_flags: str = "") -> None
         ]
     )
     host.run_cmd(
-        f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v24.09 {git_clone_flags}"
+        f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v25.02 {git_clone_flags}"
     )
 
     host.run_cmd(f"cd ComputeLibrary && scons Werror=1 -j8 {acl_build_flags}")
@@ -358,7 +358,7 @@ def checkout_repo(
     branch: str = "main",
     url: str,
     git_clone_flags: str,
-    mapping: Dict[str, Tuple[str, str]],
+    mapping: dict[str, tuple[str, str]],
 ) -> Optional[str]:
     for prefix in mapping:
         if not branch.startswith(prefix):
@@ -619,9 +619,11 @@ def build_torchaudio(
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
 
-    host.run_cmd(f"cd audio && export FFMPEG_ROOT=$(pwd)/third_party/ffmpeg && export USE_FFMPEG=1 \
+    host.run_cmd(
+        f"cd audio && export FFMPEG_ROOT=$(pwd)/third_party/ffmpeg && export USE_FFMPEG=1 \
         && ./packaging/ffmpeg/build.sh \
-        && {build_vars} python3 setup.py bdist_wheel")
+        && {build_vars} python3 setup.py bdist_wheel"
+    )
 
     wheel_name = host.list_dir("audio/dist")[0]
     embed_libgomp(host, use_conda, os.path.join("audio", "dist", wheel_name))
@@ -655,18 +657,6 @@ def configure_system(
             "sudo apt-get install -y python3-dev python3-yaml python3-setuptools python3-wheel python3-pip"
         )
     host.run_cmd("pip3 install dataclasses typing-extensions")
-    # Install and switch to gcc-8 on Ubuntu-18.04
-    if not host.using_docker() and host.ami == ubuntu18_04_ami and compiler == "gcc-8":
-        host.run_cmd("sudo apt-get install -y g++-8 gfortran-8")
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-8 100"
-        )
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-8 100"
-        )
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-8 100"
-        )
     if not use_conda:
         print("Installing Cython + numpy from PyPy")
         host.run_cmd("sudo pip3 install Cython")
@@ -679,7 +669,7 @@ def build_domains(
     branch: str = "main",
     use_conda: bool = True,
     git_clone_flags: str = "",
-) -> Tuple[str, str, str, str]:
+) -> tuple[str, str, str, str]:
     vision_wheel_name = build_torchvision(
         host, branch=branch, use_conda=use_conda, git_clone_flags=git_clone_flags
     )
@@ -706,7 +696,7 @@ def start_build(
     pytorch_build_number: Optional[str] = None,
     shallow_clone: bool = True,
     enable_mkldnn: bool = False,
-) -> Tuple[str, str, str, str, str]:
+) -> tuple[str, str, str, str, str]:
     git_clone_flags = " --depth 1 --shallow-submodules" if shallow_clone else ""
     if host.using_docker() and not use_conda:
         print("Auto-selecting conda option for docker images")
@@ -757,7 +747,7 @@ def start_build(
         version = host.check_output("cat pytorch/version.txt").strip()[:-2]
         build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1"
     if branch.startswith(("v1.", "v2.")):
-        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
+        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1 : branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
     if enable_mkldnn:
@@ -930,9 +920,9 @@ def parse_arguments():
     parser.add_argument("--debug", action="store_true")
     parser.add_argument("--build-only", action="store_true")
     parser.add_argument("--test-only", type=str)
-    parser.add_argument(
-        "--os", type=str, choices=list(os_amis.keys()), default="ubuntu20_04"
-    )
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("--os", type=str, choices=list(os_amis.keys()))
+    group.add_argument("--ami", type=str)
     parser.add_argument(
         "--python-version",
         type=str,
@@ -962,7 +952,13 @@ def parse_arguments():
 
 if __name__ == "__main__":
     args = parse_arguments()
-    ami = os_amis[args.os]
+    ami = (
+        args.ami
+        if args.ami is not None
+        else os_amis[args.os]
+        if args.os is not None
+        else ubuntu20_04_ami
+    )
     keyfile_path, key_name = compute_keyfile_path(args.key_name)
 
     if args.list_instances:
@@ -1016,7 +1012,7 @@ def parse_arguments():
         install_condaforge_python(host, args.python_version)
         sys.exit(0)
 
-    python_version = args.python_version if args.python_version is not None else "3.8"
+    python_version = args.python_version if args.python_version is not None else "3.9"
 
     if args.use_torch_from_pypi:
         configure_system(host, compiler=args.compiler, python_version=python_version)