build_aarch64_wheel.py

#!/usr/bin/env python3

import boto3
import os
import subprocess
import sys
import time
from typing import Tuple


# AMI images for us-east-1, change the following based on your ~/.aws/config
ubuntu18_04_ami = "ami-0f2b111fdc1647918"
ubuntu20_04_ami = "ami-0ea142bd244023692"


def compute_keyfile_path(key_name=None):
    if key_name is None:
        key_name = os.getenv("AWS_KEY_NAME")
    homedir_path = os.path.expanduser("~")
    default_path = os.path.join(homedir_path, ".ssh", f"{key_name}.pem")
    return os.getenv("SSH_KEY_PATH", default_path), key_name


keyfile_path = None


ec2 = boto3.resource("ec2")


def ec2_get_instances(filter_name, filter_value):
    return ec2.instances.filter(Filters=[{'Name': filter_name, 'Values': [filter_value]}])


def ec2_instances_of_type(instance_type='t4g.2xlarge'):
    return ec2_get_instances('instance-type', instance_type)


def ec2_instances_by_id(instance_id):
    rc = list(ec2_get_instances('instance-id', instance_id))
    return rc[0] if len(rc) > 0 else None


def start_instance(key_name, ami=ubuntu18_04_ami, instance_type='t4g.2xlarge'):
    inst = ec2.create_instances(ImageId=ami,
                                InstanceType=instance_type,
                                SecurityGroups=['ssh-allworld'],
                                KeyName=key_name,
                                MinCount=1,
                                MaxCount=1)[0]
    print(f'Create instance {inst.id}')
    inst.wait_until_running()
    running_inst = ec2_instances_by_id(inst.id)
    print(f'Instance started at {running_inst.public_dns_name}')
    return running_inst


def _gen_ssh_prefix(addr):
    return ["ssh", "-o", "StrictHostKeyChecking=no", "-i", keyfile_path, f"ubuntu@{addr}", "--"]


def run_ssh(addr, args):
    subprocess.check_call(_gen_ssh_prefix(addr) + (args.split() if isinstance(args, str) else args))


def check_output(addr, args):
    return subprocess.check_output(_gen_ssh_prefix(addr) + (args.split() if isinstance(args, str) else args)).decode("utf-8")


def list_dir(addr, path):
    return check_output(addr, ["ls", "-1", path]).split("\n")


def wait_for_connection(addr, port, timeout=5, attempt_cnt=5):
    import socket
    for i in range(attempt_cnt):
        try:
            with socket.create_connection((addr, port), timeout=timeout):
                return
        except (ConnectionRefusedError, socket.timeout):
            if i == attempt_cnt - 1:
                raise
            time.sleep(timeout)


def update_apt_repo(addr):
    time.sleep(5)
    run_ssh(addr, "sudo systemctl stop apt-daily.service || true")
    run_ssh(addr, "sudo systemctl stop unattended-upgrades.service || true")
    run_ssh(addr, "while systemctl is-active --quiet apt-daily.service; do sleep 1; done")
    run_ssh(addr, "while systemctl is-active --quiet unattended-upgrades.service; do sleep 1; done")
    run_ssh(addr, "sudo apt-get update")
    time.sleep(3)
    run_ssh(addr, "sudo apt-get update")


def install_condaforge(addr):
    print('Install conda-forge')
    run_ssh(addr, "curl -OL https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh")
    run_ssh(addr, "sh -f Miniforge3-Linux-aarch64.sh -b")
    run_ssh(addr, ['sed', '-i', '\'/^# If not running interactively.*/i PATH=$HOME/miniforge3/bin:$PATH\'', '.bashrc'])


def embed_libgomp(addr, use_conda, wheel_name):
    run_ssh(addr, "pip3 install auditwheel")
    run_ssh(addr, "conda install -y patchelf" if use_conda else "sudo apt-get install -y patchelf")
    from tempfile import NamedTemporaryFile
    with NamedTemporaryFile() as tmp:
        tmp.write(embed_library_script.encode('utf-8'))
        tmp.flush()
        subprocess.check_call(["scp", "-i", keyfile_path, tmp.name, f"ubuntu@{addr}:embed_library.py"])

    print('Embedding libgomp into wheel')
    run_ssh(addr, f"python3 embed_library.py {wheel_name}")


def start_build(key_name, *,
                ami=ubuntu18_04_ami,
                branch="master",
                use_conda=True,
                python_version="3.8",
                keep_running=False,
                shallow_clone=True) -> Tuple[str, str]:
    inst = start_instance(key_name, ami=ami)
    addr = inst.public_dns_name
    wait_for_connection(addr, 22)
    if use_conda:
        install_condaforge(addr)
        run_ssh(addr, f"conda install -y python={python_version} numpy pyyaml")
    git_clone_flags = " --depth 1 --shallow-submodules" if shallow_clone else ""
    print('Configuring the system')
    update_apt_repo(addr)

    run_ssh(addr, "sudo apt-get install -y ninja-build g++ git cmake gfortran unzip")
    if not use_conda:
        run_ssh(addr, "sudo apt-get install -y python3-dev python3-yaml python3-setuptools python3-wheel python3-pip")
    run_ssh(addr, "pip3 install dataclasses")
    # Install and switch to gcc-8 on Ubuntu-18.04
    if ami == ubuntu18_04_ami:
        run_ssh(addr, "sudo apt-get install -y g++-8 gfortran-8")
        run_ssh(addr, "sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-8 100")
        run_ssh(addr, "sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-8 100")
        run_ssh(addr, "sudo update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-8 100")
    if not use_conda:
        print("Installing Cython + numpy from PyPy")
        run_ssh(addr, "sudo pip3 install Cython")
        run_ssh(addr, "sudo pip3 install numpy")
    # Build OpenBLAS
    print('Building OpenBLAS')
    run_ssh(addr, f"git clone https://github.com/xianyi/OpenBLAS -b v0.3.10 {git_clone_flags}")
    # TODO: Build with USE_OPENMP=1 support
    run_ssh(addr, "pushd OpenBLAS; make NO_SHARED=1 -j8; sudo make NO_SHARED=1 install; popd")

    # Build FFTW
    print("Building FFTW3")
    run_ssh(addr, "sudo apt-get install -y ocaml ocamlbuild autoconf automake indent libtool fig2dev texinfo")
    # TODO: fix a version to build
    # TODO: consider adding flags --host=arm-linux-gnueabi --enable-single --enable-neon CC=arm-linux-gnueabi-gcc -march=armv7-a -mfloat-abi=softfp
    run_ssh(addr, f"git clone https://github.com/FFTW/fftw3 {git_clone_flags}")
    run_ssh(addr, "pushd fftw3; sh bootstrap.sh; make -j8; sudo make install; popd")

    print('Checking out PyTorch repo')
    run_ssh(addr, f"git clone --recurse-submodules -b {branch} https://github.com/pytorch/pytorch {git_clone_flags}")
    print('Building PyTorch wheel')
    build_vars=""
    if branch == 'nightly':
        build_date = check_output(addr, "cd pytorch ; git log --pretty=format:%s -1").strip().split()[0].replace("-", "")
        version = check_output(addr, "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."):
        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
    run_ssh(addr, f"cd pytorch ; {build_vars} python3 setup.py bdist_wheel")
    pytorch_wheel_name = list_dir(addr, "pytorch/dist")[0]
    embed_libgomp(addr, use_conda, os.path.join('pytorch', 'dist', pytorch_wheel_name))
    print('Copying the wheel')
    subprocess.check_call(["scp", "-i", keyfile_path, f"ubuntu@{addr}:pytorch/dist/*.whl", "."])

    print('Checking out TorchVision repo')
    if branch.startswith("v1.7.1"):
        run_ssh(addr, f"git clone https://github.com/pytorch/vision -b v0.8.2-rc2 {git_clone_flags}")
    else:
        run_ssh(addr, f"git clone https://github.com/pytorch/vision {git_clone_flags}")
    print('Installing PyTorch wheel')
    run_ssh(addr, f"pip3 install pytorch/dist/{pytorch_wheel_name}")
    print('Building TorchVision wheel')
    build_vars=""
    if branch == 'nightly':
        version = check_output(addr, ["grep", "\"version = '\"", "vision/setup.py"]).strip().split("'")[1][:-2]
        build_vars += f"BUILD_VERSION={version}.dev{build_date}"
    if branch.startswith("v1.7.1"):
        build_vars += f"BUILD_VERSION=0.8.2"

    run_ssh(addr, f"cd vision; {build_vars} python3 setup.py bdist_wheel")
    vision_wheel_name = list_dir(addr, "vision/dist")[0]
    print('Copying TorchVision wheel')
    subprocess.check_call(["scp", "-i", keyfile_path, f"ubuntu@{addr}:vision/dist/*.whl", "."])

    if keep_running:
        return pytorch_wheel_name, vision_wheel_name

    print(f'Waiting for instance {inst.id} to terminate')
    inst.terminate()
    inst.wait_until_terminated()
    return pytorch_wheel_name, vision_wheel_name


embed_library_script = """
#!/usr/bin/env python3

from auditwheel.patcher import Patchelf
from auditwheel.wheeltools import InWheelCtx
from auditwheel.elfutils import elf_file_filter
from auditwheel.repair import copylib
from auditwheel.lddtree import lddtree
import os
import shutil
import sys
from tempfile import TemporaryDirectory


def embed_library(whl_path, lib_soname):
    patcher = Patchelf()
    out_dir = TemporaryDirectory()
    whl_name = os.path.basename(whl_path)
    tmp_whl_name = os.path.join(out_dir.name, whl_name)
    with InWheelCtx(whl_path) as ctx:
        torchlib_path = os.path.join(ctx._tmpdir.name, 'torch', 'lib')
        ctx.out_wheel=tmp_whl_name
        new_lib_path, new_lib_soname = None, None
        for filename, elf in elf_file_filter(ctx.iter_files()):
            if not filename.startswith('torch/lib'):
                continue
            libtree = lddtree(filename)
            if lib_soname not in libtree['needed']:
                continue
            lib_path = libtree['libs'][lib_soname]['path']
            if lib_path is None:
                print(f"Can't embed {lib_soname} as it could not be found")
                break
            if lib_path.startswith(torchlib_path):
                continue

            if new_lib_path is None:
                new_lib_soname, new_lib_path = copylib(lib_path, torchlib_path, patcher)
            patcher.replace_needed(filename, lib_soname, new_lib_soname)
            print(f'Replacing {lib_soname} with {new_lib_soname} for {filename}')
    shutil.move(tmp_whl_name, whl_path)


if __name__ == '__main__':
    whl_name='torch-1.8.0.dev20201108-cp38-cp38-linux_aarch64.whl'
    embed_library(sys.argv[1], 'libgomp.so.1')
"""


def run_tests(ami, whl, branch='master'):
    inst = start_instance(ami)
    addr = inst.public_dns_name
    wait_for_connection(addr, 22)
    print(f'Configuring the system')
    update_apt_repo(addr)
    run_ssh(addr, "sudo apt-get install -y python3-pip git")
    run_ssh(addr, "sudo pip3 install Cython")
    run_ssh(addr, "sudo pip3 install numpy")
    subprocess.check_call(["scp", "-i", keyfile_path, whl, f"ubuntu@{addr}:"])
    run_ssh(addr, f"sudo pip3 install {whl}")
    run_ssh(addr, "python3 -c 'import torch;print(torch.rand((3,3))'")
    run_ssh(addr, f"git clone -b {branch} https://github.com/pytorch/pytorch")
    run_ssh(addr, "cd pytorch/test; python3 test_torch.py -v")


def list_instances(instance_type: str) -> None:
    print(f"All instances of type {instance_type}")
    for instance in ec2_instances_of_type(instance_type):
        print(f"{instance.id} {instance.public_dns_name} {instance.state['Name']}")


def terminate_instances(instance_type: str) -> None:
    print(f"Terminating all instances of type {instance_type}")
    instances = list(ec2_instances_of_type(instance_type))
    for instance in instances:
        print(f"Terminating {instance.id}")
        instance.terminate()
    print(f"Waiting for termination to complete")
    for instance in instances:
        instance.wait_until_terminated()


def parse_arguments():
    from argparse import ArgumentParser
    parser = ArgumentParser("Builid and test AARCH64 wheels using EC2")
    parser.add_argument("--key-name", type=str)
    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=['ubuntu18_04', 'ubuntu20_04'], default='ubuntu18_04')
    parser.add_argument("--python-version", type=str, choices=['3.6', '3.7', '3.8', '3.9'], default=None)
    parser.add_argument("--alloc-instance", action="store_true")
    parser.add_argument("--list-instances", action="store_true")
    parser.add_argument("--keep-running", action="store_true")
    parser.add_argument("--terminate-instances", action="store_true")
    parser.add_argument("--instance-type", type=str, default="t4g.2xlarge")
    parser.add_argument("--branch", type=str, default="master")
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()
    ami = ubuntu20_04_ami if args.os == 'ubuntu20_04' else ubuntu18_04_ami
    keyfile_path, key_name = compute_keyfile_path(args.key_name)

    if args.list_instances:
        list_instances(args.instance_type)
        sys.exit(0)

    if args.terminate_instances:
        terminate_instances(args.instance_type)
        sys.exit(0)

    if key_name is None:
        raise Exception("""
            Cannot start build without key_name, please specify
            --key-name argument or AWS_KEY_NAME environment variable.""")
    if keyfile_path is None or not os.path.exists(keyfile_path):
        raise Exception(f"""
            Cannot find keyfile with name: [{key_name}] in path: [{keyfile_path}], please
            check `~/.ssh/` folder or manually set SSH_KEY_PATH environment variable.""")

    if args.test_only:
        run_tests(ami, args.test_only)
        sys.exit(0)

    if args.alloc_instance:
        inst = start_instance(key_name, ami, args.instance_type)
        if args.python_version is None:
            sys.exit(0)
        addr = inst.public_dns_name
        wait_for_connection(addr, 22)
        install_condaforge(addr)
        run_ssh(addr, f"conda install -y python={args.python_version} numpy pyyaml")
        sys.exit(0)

    python_version = args.python_version if args.python_version is not None else '3.8'
    start_build(key_name, ami=ami, branch=args.branch, python_version=python_version, keep_running=args.keep_running)