Skip to content

Commit

Permalink
[aarch64] Adding scripts for aarch64 CI
Browse files Browse the repository at this point in the history 
cherrypicked:
pytorch#1293
pytorch#1362
pytorch@0157db5
pytorch@85075c1
pytorch@a007fbc
  • Loading branch information
@snadampal
snadampal committed Apr 4, 2023
1 parent 9f0fad9 commit 9317c70
Show file tree
Hide file tree
Showing 5 changed files with 561 additions and 76 deletions.
19 changes: 19 additions & 0 deletions aarch64_linux/README.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
# Aarch64 (ARM/Graviton) Support Scripts
Scripts for building aarch64 PyTorch PIP Wheels. These scripts build the following wheels:
* torch
* torchvision
* torchaudio
* torchtext
* torchdata
## Aarch64_ci_build.sh
This script is design to support CD operations within PyPi manylinux aarch64 container, and be executed in the container. It prepares the container and then executes __aarch64_wheel_ci_build.py__ to build the wheels. The script "assumes" the PyTorch repo is located at: ```/pytorch``` and will put the wheels into ```/artifacts```.
### Usage
```DESIRED_PYTHON=<PythonVersion> aarch64_ci_build.sh```

__NOTE:__ CI build is currently __EXPERMINTAL__

## Build_aarch64_wheel.py
This app allows a person to build using AWS EC3 resources and requires AWS-CLI and Boto3 with AWS credentials to support building EC2 instances for the wheel builds. Can be used in a codebuild CD or from a local system.

### Usage
```build_aarch64_wheel.py --key-name <YourPemKey> --use-docker --python 3.8 --branch <RCtag>```
52 changes: 52 additions & 0 deletions aarch64_linux/aarch64_ci_build.sh
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
#!/bin/bash
set -eux -o pipefail

# This script is used to prepare the Docker container for aarch64_ci_wheel_build.py python script
# as we need to install conda and setup the python version for the build.

CONDA_PYTHON_EXE=/opt/conda/bin/python
CONDA_EXE=/opt/conda/bin/conda
PATH=/opt/conda/bin:$PATH

###############################################################################
# Install OS dependent packages
###############################################################################
yum -y install epel-release
yum -y install less zstd

###############################################################################
# Install conda
# disable SSL_verify due to getting "Could not find a suitable TLS CA certificate bundle, invalid path"
# when using Python version, less than the conda latest
###############################################################################
echo 'Installing conda-forge'
curl -L -o /mambaforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh
chmod +x /mambaforge.sh
/mambaforge.sh -b -p /opt/conda
rm /mambaforge.sh
/opt/conda/bin/conda config --set ssl_verify False
/opt/conda/bin/conda install -y -c conda-forge python=${DESIRED_PYTHON} numpy pyyaml setuptools patchelf
python --version
conda --version

###############################################################################
# Exec libglfortran.a hack
#
# libgfortran.a from quay.io/pypa/manylinux2014_aarch64 is not compiled with -fPIC.
# This causes __stack_chk_guard@@GLIBC_2.17 on pytorch build. To solve, get
# ubuntu's libgfortran.a which is compiled with -fPIC
###############################################################################
cd ~/
curl -L -o ~/libgfortran-10-dev.deb http://ports.ubuntu.com/ubuntu-ports/pool/universe/g/gcc-10/libgfortran-10-dev_10.4.0-6ubuntu1_arm64.deb
ar x ~/libgfortran-10-dev.deb
tar --use-compress-program=unzstd -xvf data.tar.zst -C ~/
cp -f ~/usr/lib/gcc/aarch64-linux-gnu/10/libgfortran.a /opt/rh/devtoolset-10/root/usr/lib/gcc/aarch64-redhat-linux/10/

###############################################################################
# Run aarch64 builder python
###############################################################################
cd /
# adding safe directory for git as the permissions will be
# on the mounted pytorch repo
git config --global --add safe.directory /pytorch
python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn
309 changes: 309 additions & 0 deletions aarch64_linux/aarch64_wheel_ci_build.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,309 @@
#!/usr/bin/env python3

import os
import subprocess
from typing import Dict, List, Optional, Tuple


''''
Helper for getting paths for Python
'''
def list_dir(path: str) -> List[str]:
return subprocess.check_output(["ls", "-1", path]).decode().split("\n")


'''
Helper to get repo branches for specific versions
'''
def checkout_repo(branch: str = "main",
url: str = "",
git_clone_flags: str = "",
mapping: Dict[str, Tuple[str, str]] = []) -> Optional[str]:
for prefix in mapping:
if not branch.startswith(prefix):
continue
tag = f"v{mapping[prefix][0]}-{mapping[prefix][1]}"
os.system(f"git clone {url} -b {tag} {git_clone_flags}")
return mapping[prefix][0]

os.system(f"git clone {url} {git_clone_flags}")
return None


'''
Using OpenBLAS with PyTorch
'''
def build_OpenBLAS(git_clone_flags: str = "") -> None:
print('Building OpenBLAS')
os.system(f"cd /; git clone https://github.com/xianyi/OpenBLAS -b v0.3.21 {git_clone_flags}")
make_flags = "NUM_THREADS=64 USE_OPENMP=1 NO_SHARED=1 DYNAMIC_ARCH=1 TARGET=ARMV8 "
os.system(f"cd OpenBLAS; make {make_flags} -j8; make {make_flags} install; cd /; rm -rf OpenBLAS")


'''
Using ArmComputeLibrary for aarch64 PyTorch
'''
def build_ArmComputeLibrary(git_clone_flags: str = "") -> None:
print('Building Arm Compute Library')
os.system("cd / && mkdir /acl")
os.system(f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v22.11 {git_clone_flags}")
os.system(f"cd ComputeLibrary; export acl_install_dir=/acl; " \
f"scons Werror=1 -j8 debug=0 neon=1 opencl=0 os=linux openmp=1 cppthreads=0 arch=armv8.2-a multi_isa=1 build=native build_dir=$acl_install_dir/build; " \
f"cp -r arm_compute $acl_install_dir; " \
f"cp -r include $acl_install_dir; " \
f"cp -r utils $acl_install_dir; " \
f"cp -r support $acl_install_dir; " \
f"cp -r src $acl_install_dir; cd /")


'''
Script to embed libgomp to the wheels
'''
def embed_libgomp(wheel_name) -> None:
print('Embedding libgomp into wheel')
os.system(f"python3 /builder/aarch64_linux/embed_library.py {wheel_name} --update-tag")


'''
Build TorchVision wheel
'''
def build_torchvision(branch: str = "main",
git_clone_flags: str = "") -> str:
print('Checking out TorchVision repo')
build_version = checkout_repo(branch=branch,
url="https://github.com/pytorch/vision",
git_clone_flags=git_clone_flags,
mapping={
"v1.7.1": ("0.8.2", "rc2"),
"v1.8.0": ("0.9.0", "rc3"),
"v1.8.1": ("0.9.1", "rc1"),
"v1.9.0": ("0.10.0", "rc1"),
"v1.10.0": ("0.11.1", "rc1"),
"v1.10.1": ("0.11.2", "rc1"),
"v1.10.2": ("0.11.3", "rc1"),
"v1.11.0": ("0.12.0", "rc1"),
"v1.12.0": ("0.13.0", "rc4"),
"v1.12.1": ("0.13.1", "rc6"),
"v1.13.0": ("0.14.0", "rc4"),
"v1.13.1": ("0.14.1", "rc2"),
"v2.0.0": ("0.15.0", "rc2"),
})
print('Building TorchVision wheel')
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
if branch == 'nightly':
version = ''
if os.path.exists('/vision/version.txt'):
version = subprocess.check_output(['cat', '/vision/version.txt']).decode().strip()
if len(version) == 0:
# In older revisions, version was embedded in setup.py
version = subprocess.check_output(['grep', 'version', 'setup.py']).decode().strip().split('\'')[1][:-2]
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/vision').decode().replace('-','')
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
elif build_version is not None:
build_vars += f"BUILD_VERSION={build_version}"

os.system(f"cd /vision; {build_vars} python3 setup.py bdist_wheel")
wheel_name = list_dir("/vision/dist")[0]
embed_libgomp(f"/vision/dist/{wheel_name}")

print('Move TorchVision wheel to artfacts')
os.system(f"mv /vision/dist/{wheel_name} /artifacts/")
return wheel_name


'''
Build TorchAudio wheel
'''
def build_torchaudio(branch: str = "main",
git_clone_flags: str = "") -> str:
print('Checking out TorchAudio repo')
git_clone_flags += " --recurse-submodules"
build_version = checkout_repo(branch=branch,
url="https://github.com/pytorch/audio",
git_clone_flags=git_clone_flags,
mapping={
"v1.9.0": ("0.9.0", "rc2"),
"v1.10.0": ("0.10.0", "rc5"),
"v1.10.1": ("0.10.1", "rc1"),
"v1.10.2": ("0.10.2", "rc1"),
"v1.11.0": ("0.11.0", "rc1"),
"v1.12.0": ("0.12.0", "rc3"),
"v1.12.1": ("0.12.1", "rc5"),
"v1.13.0": ("0.13.0", "rc4"),
"v1.13.1": ("0.13.1", "rc2"),
"v2.0.0": ("2.0.0", "rc2"),
})
print('Building TorchAudio wheel')
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
if branch == 'nightly':
version = ''
if os.path.exists('/audio/version.txt'):
version = subprocess.check_output(['cat', '/audio/version.txt']).decode().strip()
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/audio').decode().replace('-','')
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
elif build_version is not None:
build_vars += f"BUILD_VERSION={build_version}"

os.system(f"cd /audio; {build_vars} python3 setup.py bdist_wheel")
wheel_name = list_dir("/audio/dist")[0]
embed_libgomp(f"/audio/dist/{wheel_name}")

print('Move TorchAudio wheel to artfacts')
os.system(f"mv /audio/dist/{wheel_name} /artifacts/")
return wheel_name


'''
Build TorchText wheel
'''
def build_torchtext(branch: str = "main",
git_clone_flags: str = "") -> str:
print('Checking out TorchText repo')
os.system(f"cd /")
git_clone_flags += " --recurse-submodules"
build_version = checkout_repo(branch=branch,
url="https://github.com/pytorch/text",
git_clone_flags=git_clone_flags,
mapping={
"v1.9.0": ("0.10.0", "rc1"),
"v1.10.0": ("0.11.0", "rc2"),
"v1.10.1": ("0.11.1", "rc1"),
"v1.10.2": ("0.11.2", "rc1"),
"v1.11.0": ("0.12.0", "rc1"),
"v1.12.0": ("0.13.0", "rc2"),
"v1.12.1": ("0.13.1", "rc5"),
"v1.13.0": ("0.14.0", "rc3"),
"v1.13.1": ("0.14.1", "rc1"),
"v2.0.0": ("0.15.0", "rc2"),
})
print('Building TorchText wheel')
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
if branch == 'nightly':
version = ''
if os.path.exists('/text/version.txt'):
version = subprocess.check_output(['cat', '/text/version.txt']).decode().strip()
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/text').decode().replace('-','')
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
elif build_version is not None:
build_vars += f"BUILD_VERSION={build_version}"

os.system(f"cd text; {build_vars} python3 setup.py bdist_wheel")
wheel_name = list_dir("/text/dist")[0]
embed_libgomp(f"/text/dist/{wheel_name}")

print('Move TorchText wheel to artfacts')
os.system(f"mv /text/dist/{wheel_name} /artifacts/")
return wheel_name


'''
Build TorchData wheel
'''
def build_torchdata(branch: str = "main",
git_clone_flags: str = "") -> str:
print('Checking out TorchData repo')
git_clone_flags += " --recurse-submodules"
build_version = checkout_repo(branch=branch,
url="https://github.com/pytorch/data",
git_clone_flags=git_clone_flags,
mapping={
"v1.11.0": ("0.3.0", "rc1"),
"v1.12.0": ("0.4.0", "rc3"),
"v1.12.1": ("0.4.1", "rc5"),
"v1.13.1": ("0.5.1", "rc2"),
"v2.0.0": ("0.6.0", "rc2"),
})
print('Building TorchData wheel')
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
if branch == 'nightly':
version = ''
if os.path.exists('/data/version.txt'):
version = subprocess.check_output(['cat', '/data/version.txt']).decode().strip()
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/data').decode().replace('-','')
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
elif build_version is not None:
build_vars += f"BUILD_VERSION={build_version}"

os.system(f"cd /data; {build_vars} python3 setup.py bdist_wheel")
wheel_name = list_dir("/data/dist")[0]
embed_libgomp(f"/data/dist/{wheel_name}")

print('Move TorchAudio wheel to artfacts')
os.system(f"mv /data/dist/{wheel_name} /artifacts/")
return wheel_name


def parse_arguments():
from argparse import ArgumentParser
parser = ArgumentParser("AARCH64 wheels python CD")
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("--enable-mkldnn", action="store_true")
return parser.parse_args()


'''
Entry Point
'''
if __name__ == '__main__':

args = parse_arguments()
enable_mkldnn = args.enable_mkldnn
os.system("cd /pytorch")
branch = subprocess.check_output("git rev-parse --abbrev-ref HEAD")

git_clone_flags = " --depth 1 --shallow-submodules"
os.system(f"conda install -y ninja scons")

print("Build and Install OpenBLAS")
build_OpenBLAS(git_clone_flags)

print('Building PyTorch wheel')
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
os.system(f"cd /pytorch; pip install -r requirements.txt")
os.system(f"pip install auditwheel")
os.system(f"python setup.py clean")

if branch == 'nightly' or branch == 'master':
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/pytorch').decode().replace('-','')
version = subprocess.check_output(['cat','version.txt'], cwd='/pytorch').decode().strip()[:-2]
build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1"
if branch.startswith("v1.") or branch.startswith("v2."):
build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
if enable_mkldnn:
build_ArmComputeLibrary(git_clone_flags)
print("build pytorch with mkldnn+acl backend")
os.system(f"export ACL_ROOT_DIR=/acl; export LD_LIBRARY_PATH=/acl/build; export ACL_LIBRARY=/acl/build")
build_vars += " USE_MKLDNN=ON USE_MKLDNN_ACL=ON"
os.system(f"cd /pytorch; {build_vars} python3 setup.py bdist_wheel")
print('Repair the wheel')
pytorch_wheel_name = list_dir("pytorch/dist")[0]
os.system(f"export LD_LIBRARY_PATH=/pytorch/build/lib:$LD_LIBRARY_PATH; auditwheel repair /pytorch/dist/{pytorch_wheel_name}")
print('replace the original wheel with the repaired one')
pytorch_repaired_wheel_name = list_dir("wheelhouse")[0]
os.system(f"cp /wheelhouse/{pytorch_repaired_wheel_name} /pytorch/dist/{pytorch_wheel_name}")
else:
print("build pytorch without mkldnn backend")
os.system(f"cd pytorch ; {build_vars} python3 setup.py bdist_wheel")

print("Deleting build folder")
os.system("cd /pytorch; rm -rf build")
pytorch_wheel_name = list_dir("/pytorch/dist")[0]
embed_libgomp(f"/pytorch/dist/{pytorch_wheel_name}")
print('Move PyTorch wheel to artfacts')
os.system(f"mv /pytorch/dist/{pytorch_wheel_name} /artifacts/")
print("Installing Pytorch wheel")
os.system(f"pip install /artifacts/{pytorch_wheel_name}")

vision_wheel_name = build_torchvision(branch=branch, git_clone_flags=git_clone_flags)
audio_wheel_name = build_torchaudio(branch=branch, git_clone_flags=git_clone_flags)
text_wheel_name = build_torchtext(branch=branch, git_clone_flags=git_clone_flags)
data_wheel_name = build_torchdata(branch=branch, git_clone_flags=git_clone_flags)

print(f"Wheels Created:\n" \
f"{pytorch_wheel_name}\n" \
f"{vision_wheel_name}\n" \
f"{audio_wheel_name}\n" \
f"{text_wheel_name}\n" \
f"{data_wheel_name}\n")
Loading

0 comments on commit 9317c70

Please sign in to comment.