Skip to content

Commit 9317c70

Browse files
snadampalsnadampal
committed
[aarch64] Adding scripts for aarch64 CI
cherrypicked: pytorch#1293 pytorch#1362 pytorch@0157db5 pytorch@85075c1 pytorch@a007fbc
1 parent 9f0fad9 commit 9317c70

File tree

5 files changed

+561
-76
lines changed

5 files changed

+561
-76
lines changed

‎aarch64_linux/README.md

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Aarch64 (ARM/Graviton) Support Scripts
2+
Scripts for building aarch64 PyTorch PIP Wheels. These scripts build the following wheels:
3+
* torch
4+
* torchvision
5+
* torchaudio
6+
* torchtext
7+
* torchdata
8+
## Aarch64_ci_build.sh
9+
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```.
10+
### Usage
11+
```DESIRED_PYTHON=<PythonVersion> aarch64_ci_build.sh```
12+
13+
__NOTE:__ CI build is currently __EXPERMINTAL__
14+
15+
## Build_aarch64_wheel.py
16+
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.
17+
18+
### Usage
19+
```build_aarch64_wheel.py --key-name <YourPemKey> --use-docker --python 3.8 --branch <RCtag>```

‎aarch64_linux/aarch64_ci_build.sh

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
#!/bin/bash
2+
set -eux -o pipefail
3+
4+
# This script is used to prepare the Docker container for aarch64_ci_wheel_build.py python script
5+
# as we need to install conda and setup the python version for the build.
6+
7+
CONDA_PYTHON_EXE=/opt/conda/bin/python
8+
CONDA_EXE=/opt/conda/bin/conda
9+
PATH=/opt/conda/bin:$PATH
10+
11+
###############################################################################
12+
# Install OS dependent packages
13+
###############################################################################
14+
yum -y install epel-release
15+
yum -y install less zstd
16+
17+
###############################################################################
18+
# Install conda
19+
# disable SSL_verify due to getting "Could not find a suitable TLS CA certificate bundle, invalid path"
20+
# when using Python version, less than the conda latest
21+
###############################################################################
22+
echo 'Installing conda-forge'
23+
curl -L -o /mambaforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh
24+
chmod +x /mambaforge.sh
25+
/mambaforge.sh -b -p /opt/conda
26+
rm /mambaforge.sh
27+
/opt/conda/bin/conda config --set ssl_verify False
28+
/opt/conda/bin/conda install -y -c conda-forge python=${DESIRED_PYTHON} numpy pyyaml setuptools patchelf
29+
python --version
30+
conda --version
31+
32+
###############################################################################
33+
# Exec libglfortran.a hack
34+
#
35+
# libgfortran.a from quay.io/pypa/manylinux2014_aarch64 is not compiled with -fPIC.
36+
# This causes __stack_chk_guard@@GLIBC_2.17 on pytorch build. To solve, get
37+
# ubuntu's libgfortran.a which is compiled with -fPIC
38+
###############################################################################
39+
cd ~/
40+
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
41+
ar x ~/libgfortran-10-dev.deb
42+
tar --use-compress-program=unzstd -xvf data.tar.zst -C ~/
43+
cp -f ~/usr/lib/gcc/aarch64-linux-gnu/10/libgfortran.a /opt/rh/devtoolset-10/root/usr/lib/gcc/aarch64-redhat-linux/10/
44+
45+
###############################################################################
46+
# Run aarch64 builder python
47+
###############################################################################
48+
cd /
49+
# adding safe directory for git as the permissions will be
50+
# on the mounted pytorch repo
51+
git config --global --add safe.directory /pytorch
52+
python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn

‎aarch64_linux/aarch64_wheel_ci_build.py

+309
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,309 @@
1+
#!/usr/bin/env python3
2+
3+
import os
4+
import subprocess
5+
from typing import Dict, List, Optional, Tuple
6+
7+
8+
''''
9+
Helper for getting paths for Python
10+
'''
11+
def list_dir(path: str) -> List[str]:
12+
return subprocess.check_output(["ls", "-1", path]).decode().split("\n")
13+
14+
15+
'''
16+
Helper to get repo branches for specific versions
17+
'''
18+
def checkout_repo(branch: str = "main",
19+
url: str = "",
20+
git_clone_flags: str = "",
21+
mapping: Dict[str, Tuple[str, str]] = []) -> Optional[str]:
22+
for prefix in mapping:
23+
if not branch.startswith(prefix):
24+
continue
25+
tag = f"v{mapping[prefix][0]}-{mapping[prefix][1]}"
26+
os.system(f"git clone {url} -b {tag} {git_clone_flags}")
27+
return mapping[prefix][0]
28+
29+
os.system(f"git clone {url} {git_clone_flags}")
30+
return None
31+
32+
33+
'''
34+
Using OpenBLAS with PyTorch
35+
'''
36+
def build_OpenBLAS(git_clone_flags: str = "") -> None:
37+
print('Building OpenBLAS')
38+
os.system(f"cd /; git clone https://github.com/xianyi/OpenBLAS -b v0.3.21 {git_clone_flags}")
39+
make_flags = "NUM_THREADS=64 USE_OPENMP=1 NO_SHARED=1 DYNAMIC_ARCH=1 TARGET=ARMV8 "
40+
os.system(f"cd OpenBLAS; make {make_flags} -j8; make {make_flags} install; cd /; rm -rf OpenBLAS")
41+
42+
43+
'''
44+
Using ArmComputeLibrary for aarch64 PyTorch
45+
'''
46+
def build_ArmComputeLibrary(git_clone_flags: str = "") -> None:
47+
print('Building Arm Compute Library')
48+
os.system("cd / && mkdir /acl")
49+
os.system(f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v22.11 {git_clone_flags}")
50+
os.system(f"cd ComputeLibrary; export acl_install_dir=/acl; " \
51+
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; " \
52+
f"cp -r arm_compute $acl_install_dir; " \
53+
f"cp -r include $acl_install_dir; " \
54+
f"cp -r utils $acl_install_dir; " \
55+
f"cp -r support $acl_install_dir; " \
56+
f"cp -r src $acl_install_dir; cd /")
57+
58+
59+
'''
60+
Script to embed libgomp to the wheels
61+
'''
62+
def embed_libgomp(wheel_name) -> None:
63+
print('Embedding libgomp into wheel')
64+
os.system(f"python3 /builder/aarch64_linux/embed_library.py {wheel_name} --update-tag")
65+
66+
67+
'''
68+
Build TorchVision wheel
69+
'''
70+
def build_torchvision(branch: str = "main",
71+
git_clone_flags: str = "") -> str:
72+
print('Checking out TorchVision repo')
73+
build_version = checkout_repo(branch=branch,
74+
url="https://github.com/pytorch/vision",
75+
git_clone_flags=git_clone_flags,
76+
mapping={
77+
"v1.7.1": ("0.8.2", "rc2"),
78+
"v1.8.0": ("0.9.0", "rc3"),
79+
"v1.8.1": ("0.9.1", "rc1"),
80+
"v1.9.0": ("0.10.0", "rc1"),
81+
"v1.10.0": ("0.11.1", "rc1"),
82+
"v1.10.1": ("0.11.2", "rc1"),
83+
"v1.10.2": ("0.11.3", "rc1"),
84+
"v1.11.0": ("0.12.0", "rc1"),
85+
"v1.12.0": ("0.13.0", "rc4"),
86+
"v1.12.1": ("0.13.1", "rc6"),
87+
"v1.13.0": ("0.14.0", "rc4"),
88+
"v1.13.1": ("0.14.1", "rc2"),
89+
"v2.0.0": ("0.15.0", "rc2"),
90+
})
91+
print('Building TorchVision wheel')
92+
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
93+
if branch == 'nightly':
94+
version = ''
95+
if os.path.exists('/vision/version.txt'):
96+
version = subprocess.check_output(['cat', '/vision/version.txt']).decode().strip()
97+
if len(version) == 0:
98+
# In older revisions, version was embedded in setup.py
99+
version = subprocess.check_output(['grep', 'version', 'setup.py']).decode().strip().split('\'')[1][:-2]
100+
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/vision').decode().replace('-','')
101+
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
102+
elif build_version is not None:
103+
build_vars += f"BUILD_VERSION={build_version}"
104+
105+
os.system(f"cd /vision; {build_vars} python3 setup.py bdist_wheel")
106+
wheel_name = list_dir("/vision/dist")[0]
107+
embed_libgomp(f"/vision/dist/{wheel_name}")
108+
109+
print('Move TorchVision wheel to artfacts')
110+
os.system(f"mv /vision/dist/{wheel_name} /artifacts/")
111+
return wheel_name
112+
113+
114+
'''
115+
Build TorchAudio wheel
116+
'''
117+
def build_torchaudio(branch: str = "main",
118+
git_clone_flags: str = "") -> str:
119+
print('Checking out TorchAudio repo')
120+
git_clone_flags += " --recurse-submodules"
121+
build_version = checkout_repo(branch=branch,
122+
url="https://github.com/pytorch/audio",
123+
git_clone_flags=git_clone_flags,
124+
mapping={
125+
"v1.9.0": ("0.9.0", "rc2"),
126+
"v1.10.0": ("0.10.0", "rc5"),
127+
"v1.10.1": ("0.10.1", "rc1"),
128+
"v1.10.2": ("0.10.2", "rc1"),
129+
"v1.11.0": ("0.11.0", "rc1"),
130+
"v1.12.0": ("0.12.0", "rc3"),
131+
"v1.12.1": ("0.12.1", "rc5"),
132+
"v1.13.0": ("0.13.0", "rc4"),
133+
"v1.13.1": ("0.13.1", "rc2"),
134+
"v2.0.0": ("2.0.0", "rc2"),
135+
})
136+
print('Building TorchAudio wheel')
137+
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
138+
if branch == 'nightly':
139+
version = ''
140+
if os.path.exists('/audio/version.txt'):
141+
version = subprocess.check_output(['cat', '/audio/version.txt']).decode().strip()
142+
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/audio').decode().replace('-','')
143+
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
144+
elif build_version is not None:
145+
build_vars += f"BUILD_VERSION={build_version}"
146+
147+
os.system(f"cd /audio; {build_vars} python3 setup.py bdist_wheel")
148+
wheel_name = list_dir("/audio/dist")[0]
149+
embed_libgomp(f"/audio/dist/{wheel_name}")
150+
151+
print('Move TorchAudio wheel to artfacts')
152+
os.system(f"mv /audio/dist/{wheel_name} /artifacts/")
153+
return wheel_name
154+
155+
156+
'''
157+
Build TorchText wheel
158+
'''
159+
def build_torchtext(branch: str = "main",
160+
git_clone_flags: str = "") -> str:
161+
print('Checking out TorchText repo')
162+
os.system(f"cd /")
163+
git_clone_flags += " --recurse-submodules"
164+
build_version = checkout_repo(branch=branch,
165+
url="https://github.com/pytorch/text",
166+
git_clone_flags=git_clone_flags,
167+
mapping={
168+
"v1.9.0": ("0.10.0", "rc1"),
169+
"v1.10.0": ("0.11.0", "rc2"),
170+
"v1.10.1": ("0.11.1", "rc1"),
171+
"v1.10.2": ("0.11.2", "rc1"),
172+
"v1.11.0": ("0.12.0", "rc1"),
173+
"v1.12.0": ("0.13.0", "rc2"),
174+
"v1.12.1": ("0.13.1", "rc5"),
175+
"v1.13.0": ("0.14.0", "rc3"),
176+
"v1.13.1": ("0.14.1", "rc1"),
177+
"v2.0.0": ("0.15.0", "rc2"),
178+
})
179+
print('Building TorchText wheel')
180+
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
181+
if branch == 'nightly':
182+
version = ''
183+
if os.path.exists('/text/version.txt'):
184+
version = subprocess.check_output(['cat', '/text/version.txt']).decode().strip()
185+
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/text').decode().replace('-','')
186+
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
187+
elif build_version is not None:
188+
build_vars += f"BUILD_VERSION={build_version}"
189+
190+
os.system(f"cd text; {build_vars} python3 setup.py bdist_wheel")
191+
wheel_name = list_dir("/text/dist")[0]
192+
embed_libgomp(f"/text/dist/{wheel_name}")
193+
194+
print('Move TorchText wheel to artfacts')
195+
os.system(f"mv /text/dist/{wheel_name} /artifacts/")
196+
return wheel_name
197+
198+
199+
'''
200+
Build TorchData wheel
201+
'''
202+
def build_torchdata(branch: str = "main",
203+
git_clone_flags: str = "") -> str:
204+
print('Checking out TorchData repo')
205+
git_clone_flags += " --recurse-submodules"
206+
build_version = checkout_repo(branch=branch,
207+
url="https://github.com/pytorch/data",
208+
git_clone_flags=git_clone_flags,
209+
mapping={
210+
"v1.11.0": ("0.3.0", "rc1"),
211+
"v1.12.0": ("0.4.0", "rc3"),
212+
"v1.12.1": ("0.4.1", "rc5"),
213+
"v1.13.1": ("0.5.1", "rc2"),
214+
"v2.0.0": ("0.6.0", "rc2"),
215+
})
216+
print('Building TorchData wheel')
217+
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
218+
if branch == 'nightly':
219+
version = ''
220+
if os.path.exists('/data/version.txt'):
221+
version = subprocess.check_output(['cat', '/data/version.txt']).decode().strip()
222+
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/data').decode().replace('-','')
223+
build_vars += f"BUILD_VERSION={version}.dev{build_date}"
224+
elif build_version is not None:
225+
build_vars += f"BUILD_VERSION={build_version}"
226+
227+
os.system(f"cd /data; {build_vars} python3 setup.py bdist_wheel")
228+
wheel_name = list_dir("/data/dist")[0]
229+
embed_libgomp(f"/data/dist/{wheel_name}")
230+
231+
print('Move TorchAudio wheel to artfacts')
232+
os.system(f"mv /data/dist/{wheel_name} /artifacts/")
233+
return wheel_name
234+
235+
236+
def parse_arguments():
237+
from argparse import ArgumentParser
238+
parser = ArgumentParser("AARCH64 wheels python CD")
239+
parser.add_argument("--debug", action="store_true")
240+
parser.add_argument("--build-only", action="store_true")
241+
parser.add_argument("--test-only", type=str)
242+
parser.add_argument("--enable-mkldnn", action="store_true")
243+
return parser.parse_args()
244+
245+
246+
'''
247+
Entry Point
248+
'''
249+
if __name__ == '__main__':
250+
251+
args = parse_arguments()
252+
enable_mkldnn = args.enable_mkldnn
253+
os.system("cd /pytorch")
254+
branch = subprocess.check_output("git rev-parse --abbrev-ref HEAD")
255+
256+
git_clone_flags = " --depth 1 --shallow-submodules"
257+
os.system(f"conda install -y ninja scons")
258+
259+
print("Build and Install OpenBLAS")
260+
build_OpenBLAS(git_clone_flags)
261+
262+
print('Building PyTorch wheel')
263+
build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
264+
os.system(f"cd /pytorch; pip install -r requirements.txt")
265+
os.system(f"pip install auditwheel")
266+
os.system(f"python setup.py clean")
267+
268+
if branch == 'nightly' or branch == 'master':
269+
build_date = subprocess.check_output(['git','log','--pretty=format:%cs','-1'], cwd='/pytorch').decode().replace('-','')
270+
version = subprocess.check_output(['cat','version.txt'], cwd='/pytorch').decode().strip()[:-2]
271+
build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1"
272+
if branch.startswith("v1.") or branch.startswith("v2."):
273+
build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
274+
if enable_mkldnn:
275+
build_ArmComputeLibrary(git_clone_flags)
276+
print("build pytorch with mkldnn+acl backend")
277+
os.system(f"export ACL_ROOT_DIR=/acl; export LD_LIBRARY_PATH=/acl/build; export ACL_LIBRARY=/acl/build")
278+
build_vars += " USE_MKLDNN=ON USE_MKLDNN_ACL=ON"
279+
os.system(f"cd /pytorch; {build_vars} python3 setup.py bdist_wheel")
280+
print('Repair the wheel')
281+
pytorch_wheel_name = list_dir("pytorch/dist")[0]
282+
os.system(f"export LD_LIBRARY_PATH=/pytorch/build/lib:$LD_LIBRARY_PATH; auditwheel repair /pytorch/dist/{pytorch_wheel_name}")
283+
print('replace the original wheel with the repaired one')
284+
pytorch_repaired_wheel_name = list_dir("wheelhouse")[0]
285+
os.system(f"cp /wheelhouse/{pytorch_repaired_wheel_name} /pytorch/dist/{pytorch_wheel_name}")
286+
else:
287+
print("build pytorch without mkldnn backend")
288+
os.system(f"cd pytorch ; {build_vars} python3 setup.py bdist_wheel")
289+
290+
print("Deleting build folder")
291+
os.system("cd /pytorch; rm -rf build")
292+
pytorch_wheel_name = list_dir("/pytorch/dist")[0]
293+
embed_libgomp(f"/pytorch/dist/{pytorch_wheel_name}")
294+
print('Move PyTorch wheel to artfacts')
295+
os.system(f"mv /pytorch/dist/{pytorch_wheel_name} /artifacts/")
296+
print("Installing Pytorch wheel")
297+
os.system(f"pip install /artifacts/{pytorch_wheel_name}")
298+
299+
vision_wheel_name = build_torchvision(branch=branch, git_clone_flags=git_clone_flags)
300+
audio_wheel_name = build_torchaudio(branch=branch, git_clone_flags=git_clone_flags)
301+
text_wheel_name = build_torchtext(branch=branch, git_clone_flags=git_clone_flags)
302+
data_wheel_name = build_torchdata(branch=branch, git_clone_flags=git_clone_flags)
303+
304+
print(f"Wheels Created:\n" \
305+
f"{pytorch_wheel_name}\n" \
306+
f"{vision_wheel_name}\n" \
307+
f"{audio_wheel_name}\n" \
308+
f"{text_wheel_name}\n" \
309+
f"{data_wheel_name}\n")

0 commit comments

Comments
 (0)