[CI] Use native arm64 worker in GHAction to build M1 wheel

.github/workflows/python_wheels.yml

@@ -5,6 +5,10 @@ on: [push, pull_request]
 permissions:
   contents: read # to fetch code (actions/checkout)
 
+defaults:
+  run:
+    shell: bash -l {0}
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
@@ -16,30 +20,32 @@ jobs:
     strategy:
       matrix:
         include:
-        - os: macos-latest
+        - os: macos-13
           platform_id: macosx_x86_64
-        - os: macos-latest
+        - os: macos-14
           platform_id: macosx_arm64
     steps:
     - uses: actions/checkout@a12a3943b4bdde767164f792f33f40b04645d846 # v3.0.0
       with:
         submodules: 'true'
-    - name: Setup Python
-      uses: actions/setup-python@0a5c61591373683505ea898e09a3ea4f39ef2b9c # v5.0.0
+    - uses: conda-incubator/[email protected]
       with:
-        python-version: "3.8"
+        miniforge-variant: Mambaforge
+        miniforge-version: latest
+        python-version: 3.9
+        use-mamba: true
     - name: Build wheels
       run: bash tests/ci_build/build_python_wheels.sh ${{ matrix.platform_id }} ${{ github.sha }}
     - name: Extract branch name
-      shell: bash
-      run: echo "##[set-output name=branch;]$(echo ${GITHUB_REF#refs/heads/})"
+      run: |
+        echo "branch=${GITHUB_REF#refs/heads/}" >> "$GITHUB_OUTPUT"
       id: extract_branch
       if: github.ref == 'refs/heads/master' || contains(github.ref, 'refs/heads/release_')
     - name: Upload Python wheel
       if: github.ref == 'refs/heads/master' || contains(github.ref, 'refs/heads/release_')
       run: |
         python -m pip install awscli
-        python -m awscli s3 cp wheelhouse/*.whl s3://xgboost-nightly-builds/${{ steps.extract_branch.outputs.branch }}/ --acl public-read
+        python -m awscli s3 cp wheelhouse/*.whl s3://xgboost-nightly-builds/${{ steps.extract_branch.outputs.branch }}/ --acl public-read --region us-west-2
       env:
         AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID_IAM_S3_UPLOADER }}
         AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY_IAM_S3_UPLOADER }}

python-package/packager/nativelib.py

@@ -53,12 +53,6 @@ def _build(*, generator: str) -> None:
         ]
         cmake_cmd.extend(build_config.get_cmake_args())
 
-        # Flag for cross-compiling for Apple Silicon
-        # We use environment variable because it's the only way to pass down custom
-        # flags through the cibuildwheel package, which calls `pip wheel` command.
-        if "CIBW_TARGET_OSX_ARM64" in os.environ:
-            cmake_cmd.append("-DCMAKE_OSX_ARCHITECTURES=arm64")
-
         logger.info("CMake args: %s", str(cmake_cmd))
         subprocess.check_call(cmake_cmd, cwd=build_dir)
 

tests/ci_build/build_python_wheels.sh

@@ -25,16 +25,16 @@ if [[ "$platform_id" == macosx_* ]]; then
         # arm64 builds must cross compile because CI is on x64
         # cibuildwheel will take care of cross-compilation.
         wheel_tag=macosx_12_0_arm64
-        cpython_ver=38
-        setup_env_var='CIBW_TARGET_OSX_ARM64=1'  # extra flag to be passed to xgboost.packager backend
-        export PYTHON_CROSSENV=1
+        cpython_ver=39
+        cibw_archs=arm64
         export MACOSX_DEPLOYMENT_TARGET=12.0
         #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
         OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hf3c4609_1-osx-arm64.tar.bz2"
     elif [[ "$platform_id" == macosx_x86_64 ]]; then
         # MacOS, Intel
         wheel_tag=macosx_10_15_x86_64.macosx_11_0_x86_64.macosx_12_0_x86_64
-        cpython_ver=38
+        cpython_ver=39
+        cibw_archs=x86_64
         export MACOSX_DEPLOYMENT_TARGET=10.15
         #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
         OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hda6cdc1_1-osx-64.tar.bz2"
@@ -44,13 +44,13 @@ if [[ "$platform_id" == macosx_* ]]; then
     fi
     # Set up environment variables to configure cibuildwheel
     export CIBW_BUILD=cp${cpython_ver}-${platform_id}
-    export CIBW_ARCHS=all
+    export CIBW_ARCHS=${cibw_archs}
     export CIBW_ENVIRONMENT=${setup_env_var}
     export CIBW_TEST_SKIP='*-macosx_arm64'
     export CIBW_BUILD_VERBOSITY=3
 
-    sudo conda create -n build $OPENMP_URL
-    PREFIX="/usr/local/miniconda/envs/build"
+    mamba create -n build $OPENMP_URL
+    PREFIX="$HOME/miniconda3/envs/build"
 
     # Set up build flags for cibuildwheel
     # This is needed to bundle libomp lib we downloaded earlier

tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py

@@ -113,8 +113,8 @@ def run_with_dask_array(DMatrixT: Type, client: Client) -> None:
     cp.cuda.runtime.setDevice(0)
     X, y, _ = generate_array()
 
-    X = X.map_blocks(cp.asarray)
-    y = y.map_blocks(cp.asarray)
+    X = X.map_blocks(cp.asarray)  # type: ignore
+    y = y.map_blocks(cp.asarray)  # type: ignore
     dtrain = DMatrixT(client, X, y)
     out = dxgb.train(
         client,
@@ -648,8 +648,8 @@ async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainRetur
         import cupy as cp
 
         X, y, _ = generate_array()
-        X = X.map_blocks(cp.array)
-        y = y.map_blocks(cp.array)
+        X = X.map_blocks(cp.array)  # type: ignore
+        y = y.map_blocks(cp.array)  # type: ignore
 
         m = await xgb.dask.DaskQuantileDMatrix(client, X, y)
         output = await xgb.dask.train(