Run tests and build Python wheels for aarch64 architecture (#3948)

* Update setup.sh

* Update test.sh

* Update test_dask.py

* Update test_engine.py

* Update .vsts-ci.yml

.ci/setup.sh

@@ -95,7 +95,12 @@ else  # Linux
             cmake
     fi
     if [[ $SETUP_CONDA != "false" ]]; then
+        ARCH=$(uname -m)
+        if [[ $ARCH == "x86_64" ]]; then
             curl -sL -o conda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+        else
+            curl -sL -o conda.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-${ARCH}.sh
+        fi
     fi
 fi
 

.ci/test.sh

@@ -127,9 +127,15 @@ elif [[ $TASK == "bdist" ]]; then
             cp dist/lightgbm-$LGB_VER-py3-none-macosx*.whl $BUILD_ARTIFACTSTAGINGDIRECTORY
         fi
     else
-        cd $BUILD_DIRECTORY/python-package && python setup.py bdist_wheel --plat-name=manylinux1_x86_64 --python-tag py3 || exit -1
+        ARCH=$(uname -m)
+        if [[ $ARCH == "x86_64" ]]; then
+            PLATFORM="manylinux1_x86_64"
+        else
+            PLATFORM="manylinux2014_$ARCH"
+        fi
+        cd $BUILD_DIRECTORY/python-package && python setup.py bdist_wheel --plat-name=$PLATFORM --python-tag py3 || exit -1
         if [[ $PRODUCES_ARTIFACTS == "true" ]]; then
-            cp dist/lightgbm-$LGB_VER-py3-none-manylinux1_x86_64.whl $BUILD_ARTIFACTSTAGINGDIRECTORY
+            cp dist/lightgbm-$LGB_VER-py3-none-$PLATFORM.whl $BUILD_ARTIFACTSTAGINGDIRECTORY
         fi
     fi
     pip install --user $BUILD_DIRECTORY/python-package/dist/*.whl || exit -1

.vsts-ci.yml

@@ -138,6 +138,68 @@ jobs:
   - bash: $(Build.SourcesDirectory)/.ci/test.sh
     displayName: Test
 ###########################################
+- job: QEMU_multiarch
+###########################################
+  variables:
+    COMPILER: gcc
+    OS_NAME: 'linux'
+    PRODUCES_ARTIFACTS: 'true'
+  pool:
+    vmImage: ubuntu-latest
+  timeoutInMinutes: 120
+  strategy:
+    matrix:
+      bdist:
+        TASK: bdist
+        ARCH: aarch64
+  steps:
+  - script: |
+      sudo apt-get update
+      sudo apt-get install --no-install-recommends -y \
+        binfmt-support \
+        qemu \
+        qemu-user \
+        qemu-user-static
+    displayName: 'Install QEMU'
+  - script: |
+      docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
+    displayName: 'Enable Docker multi-architecture support'
+  - script: |
+      export ROOT_DOCKER_FOLDER=/LightGBM
+      cat > docker.env <<EOF
+      AZURE=$AZURE
+      OS_NAME=$OS_NAME
+      COMPILER=$COMPILER
+      TASK=$TASK
+      METHOD=$METHOD
+      CONDA_ENV=$CONDA_ENV
+      PYTHON_VERSION=$PYTHON_VERSION
+      BUILD_DIRECTORY=$ROOT_DOCKER_FOLDER
+      LGB_VER=$(head -n 1 VERSION.txt)
+      PRODUCES_ARTIFACTS=$PRODUCES_ARTIFACTS
+      BUILD_ARTIFACTSTAGINGDIRECTORY=$BUILD_ARTIFACTSTAGINGDIRECTORY
+      EOF
+      cat > docker-script.sh <<EOF
+      export CONDA=\$HOME/miniconda
+      export PATH=\$CONDA/bin:\$PATH
+      $ROOT_DOCKER_FOLDER/.ci/setup.sh || exit -1
+      $ROOT_DOCKER_FOLDER/.ci/test.sh || exit -1
+      EOF
+      docker run \
+        --rm \
+        --env-file docker.env \
+        -v "$(Build.SourcesDirectory)":"$ROOT_DOCKER_FOLDER" \
+        -v "$(Build.ArtifactStagingDirectory)":"$(Build.ArtifactStagingDirectory)" \
+        "quay.io/pypa/manylinux2014_$ARCH" \
+        /bin/bash $ROOT_DOCKER_FOLDER/docker-script.sh
+    displayName: 'Setup and run tests'
+  - task: PublishBuildArtifacts@1
+    condition: and(succeeded(), in(variables['TASK'], 'bdist'), not(startsWith(variables['Build.SourceBranch'], 'refs/pull/')))
+    inputs:
+      pathtoPublish: '$(Build.ArtifactStagingDirectory)'
+      artifactName: PackageAssets
+      artifactType: container
+###########################################
 - job: MacOS
 ###########################################
   variables:
@@ -219,6 +281,7 @@ jobs:
   dependsOn:
   - Linux
   - Linux_latest
+  - QEMU_multiarch
   - MacOS
   - Windows
   condition: and(succeeded(), not(startsWith(variables['Build.SourceBranch'], 'refs/pull/')))

tests/python_package_test/test_dask.py

@@ -5,6 +5,7 @@ import inspect
 import pickle
 import socket
 from itertools import groupby
+from platform import machine
 from os import getenv
 from sys import platform
 
@@ -43,7 +44,8 @@ group_sizes = [5, 5, 5, 10, 10, 10, 20, 20, 20, 50, 50]
 
 pytestmark = [
     pytest.mark.skipif(getenv('TASK', '') == 'mpi', reason='Fails to run with MPI interface'),
-    pytest.mark.skipif(getenv('TASK', '') == 'gpu', reason='Fails to run with GPU interface')
+    pytest.mark.skipif(getenv('TASK', '') == 'gpu', reason='Fails to run with GPU interface'),
+    pytest.mark.skipif(machine() != 'x86_64', reason='Fails to run with non-x86_64 architecture')
 ]
 
 

tests/python_package_test/test_engine.py

@@ -4,6 +4,7 @@ import itertools
 import math
 import os
 import pickle
+import platform
 import random
 
 import numpy as np
@@ -1044,6 +1045,9 @@ def test_contribs_sparse():
     # convert data to dense and get back same contribs
     contribs_dense = gbm.predict(X_test.toarray(), pred_contrib=True)
     # validate the values are the same
+    if platform.machine() == 'aarch64':
+        np.testing.assert_allclose(contribs_csr.toarray(), contribs_dense, rtol=1, atol=1e-12)
+    else:
         np.testing.assert_allclose(contribs_csr.toarray(), contribs_dense)
     assert (np.linalg.norm(gbm.predict(X_test, raw_score=True)
                            - np.sum(contribs_dense, axis=1)) < 1e-4)
@@ -1052,6 +1056,9 @@ def test_contribs_sparse():
     contribs_csc = gbm.predict(X_test_csc, pred_contrib=True)
     assert isspmatrix_csc(contribs_csc)
     # validate the values are the same
+    if platform.machine() == 'aarch64':
+        np.testing.assert_allclose(contribs_csc.toarray(), contribs_dense, rtol=1, atol=1e-12)
+    else:
         np.testing.assert_allclose(contribs_csc.toarray(), contribs_dense)
 
 
@@ -1084,6 +1091,9 @@ def test_contribs_sparse_multiclass():
     contribs_csr_array = np.swapaxes(np.array([sparse_array.todense() for sparse_array in contribs_csr]), 0, 1)
     contribs_csr_arr_re = contribs_csr_array.reshape((contribs_csr_array.shape[0],
                                                       contribs_csr_array.shape[1] * contribs_csr_array.shape[2]))
+    if platform.machine() == 'aarch64':
+        np.testing.assert_allclose(contribs_csr_arr_re, contribs_dense, rtol=1, atol=1e-12)
+    else:
         np.testing.assert_allclose(contribs_csr_arr_re, contribs_dense)
     contribs_dense_re = contribs_dense.reshape(contribs_csr_array.shape)
     assert np.linalg.norm(gbm.predict(X_test, raw_score=True) - np.sum(contribs_dense_re, axis=2)) < 1e-4
@@ -1097,6 +1107,9 @@ def test_contribs_sparse_multiclass():
     contribs_csc_array = np.swapaxes(np.array([sparse_array.todense() for sparse_array in contribs_csc]), 0, 1)
     contribs_csc_array = contribs_csc_array.reshape((contribs_csc_array.shape[0],
                                                      contribs_csc_array.shape[1] * contribs_csc_array.shape[2]))
+    if platform.machine() == 'aarch64':
+        np.testing.assert_allclose(contribs_csc_array, contribs_dense, rtol=1, atol=1e-12)
+    else:
         np.testing.assert_allclose(contribs_csc_array, contribs_dense)