From 21334e6e2012b1ab86b033726bc03cd7ce5907fb Mon Sep 17 00:00:00 2001
From: Matt Davis <daviswm@amazon.com>
Date: Thu, 1 Oct 2020 14:33:57 -0500
Subject: [PATCH] Add arm64 wheel to travis-ci

---
 .ci/setup.sh | 11 ++++++++++-
 .ci/test.sh  | 11 +++++++++--
 .travis.yml  | 29 ++++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/.ci/setup.sh b/.ci/setup.sh
index 62f34523f699..3e5e7da7e0a2 100755
--- a/.ci/setup.sh
+++ b/.ci/setup.sh
@@ -19,6 +19,11 @@ if [[ $OS_NAME == "macos" ]]; then
     fi
     wget -q -O conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
 else  # Linux
+    if [[ $(uname -m) == "aarch64" ]]; then
+        sudo rm -f $(which cmake)  # remove x86 cmake
+        sudo apt-get update
+        sudo apt-get install -y cmake
+    fi
     if [[ $TASK == "mpi" ]]; then
         sudo apt-get update
         sudo apt-get install --no-install-recommends -y libopenmpi-dev openmpi-bin
@@ -37,7 +42,11 @@ else  # Linux
         echo libamdocl64.so > $OPENCL_VENDOR_PATH/amdocl64.icd
     fi
     if [[ $TRAVIS == "true" ]] || [[ $GITHUB_ACTIONS == "true" ]]; then
-        wget -q -O conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+        CONDA_SCRIPT=https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+        if [[ $(uname -m) == "aarch64" ]]; then
+            CONDA_SCRIPT=https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh
+        fi
+        wget -q -O conda.sh $CONDA_SCRIPT
     fi
 fi
 
diff --git a/.ci/test.sh b/.ci/test.sh
index 392b48e2094b..f783440555f6 100755
--- a/.ci/test.sh
+++ b/.ci/test.sh
@@ -106,9 +106,16 @@ elif [[ $TASK == "bdist" ]]; then
             cp dist/lightgbm-$LGB_VER-py2.py3-none-macosx*.whl $BUILD_ARTIFACTSTAGINGDIRECTORY
         fi
     else
-        cd $BUILD_DIRECTORY/python-package && python setup.py bdist_wheel --plat-name=manylinux1_x86_64 --universal || exit -1
+        cd $BUILD_DIRECTORY/python-package
+        if [[ $(uname -m) == "aarch64" ]]; then
+            python setup.py bdist_wheel --plat-name=manylinux2014_aarch64 || exit -1
+            WHEEL=dist/lightgbm-$LGB_VER-py3-none-manylinux2014_aarch64.whl
+        else
+            python setup.py bdist_wheel --plat-name=manylinux1_x86_64 --universal || exit -1
+            WHEEL=dist/lightgbm-$LGB_VER-py2.py3-none-manylinux1_x86_64.whl
+        fi
         if [[ $AZURE == "true" ]]; then
-            cp dist/lightgbm-$LGB_VER-py2.py3-none-manylinux1_x86_64.whl $BUILD_ARTIFACTSTAGINGDIRECTORY
+            cp $WHEEL $BUILD_ARTIFACTSTAGINGDIRECTORY
         fi
     fi
     pip install --user $BUILD_DIRECTORY/python-package/dist/*.whl || exit -1
diff --git a/.travis.yml b/.travis.yml
index 0b4dd6bef5ef..0b6f1db25bc8 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -6,11 +6,17 @@ git:
   submodules: true
   depth: 5
 
+
+arch:
+  - amd64
+  - arm64-graviton2
 os:
   - linux
   - osx
 dist: focal
 osx_image: xcode12
+virt: vm
+group: edge
 
 env:
   global:  # default values
@@ -37,6 +43,20 @@ matrix:
       env: TASK=lint
     - os: osx
       env: TASK=check-docs
+    - arch: arm64-graviton2
+      env: TASK=sdist PYTHON_VERSION=2.7
+    - arch: arm64-graviton2
+      env: TASK=gpu METHOD=source
+    - arch: arm64-graviton2
+      env: TASK=gpu METHOD=pip PYTHON_VERSION=3.6
+    - arch: arm64-graviton2
+      env: TASK=lint
+    - arch: arm64-graviton2
+      env: TASK=check-docs
+  include:
+    - arch: arm64-graviton2
+      os: linux
+      env: TASK=sdist PYTHON_VERSION=3.6
 
 before_install:
   - test -n $CC  && unset CC
@@ -45,11 +65,18 @@ before_install:
   - if [[ $TRAVIS_OS_NAME == "osx" ]]; then
         export OS_NAME="macos";
         export COMPILER="gcc";
+    elif [[ $(uname -m) == "aarch64" ]]; then
+        export OS_NAME="linux";
+        export COMPILER="gcc";
     else
         export OS_NAME="linux";
         export COMPILER="clang";
     fi
-  - export CONDA="$HOME/miniconda"
+  - if [[ $(uname -m) == "aarch64" ]]; then
+        export CONDA="$HOME/miniforge3";
+    else
+        export CONDA="$HOME/miniconda";
+    fi
   - export PATH="$CONDA/bin:$PATH"
   - export CONDA_ENV="test-env"
   - export LGB_VER=$(head -n 1 VERSION.txt)