[backport] Build xgboost-cpu for manylinux_2_28_x86_64 (dmlc#11406) (dmlc#11407)

Author: hcho3

.github/workflows/main.yml

@@ -36,28 +36,6 @@ jobs:
             --prefix cache/${{ github.run_id }}/build-cpu \
             ./xgboost
 
-  build-cpu-arm64:
-    name: Build CPU ARM64 + manylinux_2_28_aarch64 wheel
-    runs-on:
-      - runs-on=${{ github.run_id }}
-      - runner=linux-arm64-cpu
-      - tag=build-cpu-arm64
-    steps:
-      # Restart Docker daemon so that it recognizes the ephemeral disks
-      - run: sudo systemctl restart docker
-      - uses: actions/checkout@v4
-        with:
-          submodules: "true"
-      - name: Log into Docker registry (AWS ECR)
-        run: bash ops/pipeline/login-docker-registry.sh
-      - run: bash ops/pipeline/build-cpu-arm64.sh
-      - name: Stash files
-        run: |
-          python3 ops/pipeline/manage-artifacts.py upload \
-            --s3-bucket ${{ env.RUNS_ON_S3_BUCKET_CACHE }} \
-            --prefix cache/${{ github.run_id }}/build-cpu-arm64 \
-            ./xgboost python-package/dist/*.whl
-
   build-cuda:
     name: Build CUDA + manylinux_2_28_x86_64 wheel
     runs-on:
@@ -122,13 +100,35 @@ jobs:
           bash ops/pipeline/build-cuda.sh \
             xgb-ci.gpu_build_rockylinux8_dev_ver enable-rmm
 
-  build-manylinux2014:
+  build-python-wheels-arm64:
+    name: Build manylinux_2_28_aarch64 wheel
+    runs-on:
+      - runs-on=${{ github.run_id }}
+      - runner=linux-arm64-cpu
+      - tag=build-python-wheels-arm64
+    steps:
+      # Restart Docker daemon so that it recognizes the ephemeral disks
+      - run: sudo systemctl restart docker
+      - uses: actions/checkout@v4
+        with:
+          submodules: "true"
+      - name: Log into Docker registry (AWS ECR)
+        run: bash ops/pipeline/login-docker-registry.sh
+      - run: bash ops/pipeline/build-python-wheels-arm64.sh
+      - name: Stash files
+        run: |
+          python3 ops/pipeline/manage-artifacts.py upload \
+            --s3-bucket ${{ env.RUNS_ON_S3_BUCKET_CACHE }} \
+            --prefix cache/${{ github.run_id }}/build-python-wheels-arm64 \
+            ./xgboost python-package/dist/*.whl
+
+  build-python-wheels-manylinux2014:
     name: Build manylinux2014_${{ matrix.arch }} wheel
     runs-on:
       - runs-on
       - runner=${{ matrix.runner }}
       - run-id=${{ github.run_id }}
-      - tag=main-build-manylinux2014-${{ matrix.arch }}
+      - tag=main-build-python-wheels-manylinux2014-${{ matrix.arch }}
     strategy:
       fail-fast: false
       matrix:
@@ -145,7 +145,42 @@ jobs:
           submodules: "true"
       - name: Log into Docker registry (AWS ECR)
         run: bash ops/pipeline/login-docker-registry.sh
-      - run: bash ops/pipeline/build-manylinux2014.sh ${{ matrix.arch }}
+      - run: bash ops/pipeline/build-python-wheels-manylinux2014.sh ${{ matrix.arch }}
+
+  build-python-wheels-cpu:
+    name: Build CPU wheel for ${{ matrix.manylinux_target }}_${{ matrix.arch }}
+    runs-on:
+      - runs-on
+      - runner=${{ matrix.runner }}
+      - run-id=${{ github.run_id }}
+      - tag=main-build-python-wheels-cpu-${{ matrix.manylinux_target }}-${{ matrix.arch }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - manylinux_target: manylinux2014
+          arch: aarch64
+          runner: linux-arm64-cpu
+        - manylinux_target: manylinux2014
+          arch: x86_64
+          runner: linux-amd64-cpu
+        - manylinux_target: manylinux_2_28
+          arch: aarch64
+          runner: linux-arm64-cpu
+        - manylinux_target: manylinux_2_28
+          arch: x86_64
+          runner: linux-amd64-cpu
+    steps:
+      # Restart Docker daemon so that it recognizes the ephemeral disks
+      - run: sudo systemctl restart docker
+      - uses: actions/checkout@v4
+        with:
+          submodules: "true"
+      - name: Log into Docker registry (AWS ECR)
+        run: bash ops/pipeline/login-docker-registry.sh
+      - run: |
+          bash ops/pipeline/build-python-wheels-cpu.sh \
+            ${{ matrix.manylinux_target }} ${{ matrix.arch }}
 
   build-gpu-rpkg:
     name: Build GPU-enabled R package
@@ -163,7 +198,6 @@ jobs:
         run: bash ops/pipeline/login-docker-registry.sh
       - run: bash ops/pipeline/build-gpu-rpkg.sh
 
-
   test-cpp-gpu:
     name: >-
       Run Google Tests with GPUs
@@ -208,7 +242,7 @@ jobs:
 
   test-python-wheel:
     name: Run Python tests (${{ matrix.description }})
-    needs: [build-cuda, build-cpu-arm64]
+    needs: [build-cuda, build-python-wheels-arm64]
     runs-on:
       - runs-on
       - runner=${{ matrix.runner }}
@@ -238,7 +272,7 @@ jobs:
             image_repo: xgb-ci.aarch64
             suite: cpu-arm64
             runner: linux-arm64-cpu
-            artifact_from: build-cpu-arm64
+            artifact_from: build-python-wheels-arm64
     steps:
       # Restart Docker daemon so that it recognizes the ephemeral disks
       - run: sudo systemctl restart docker

ops/pipeline/build-cpu-arm64-impl.sh renamed to ops/pipeline/build-python-wheels-arm64-impl.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-## Build and test XGBoost with ARM64 CPU
+## Build and test XGBoost with ARM64 CPU (no GPU, no federated learning)
 
 set -euo pipefail
 
@@ -14,13 +14,15 @@ source ops/pipeline/get-docker-registry-details.sh
 source ops/pipeline/get-image-tag.sh
 
 WHEEL_TAG=manylinux_2_28_aarch64
-IMAGE_URI=${DOCKER_REGISTRY_URL}/xgb-ci.aarch64:${IMAGE_TAG}
+IMAGE_URI=${DOCKER_REGISTRY_URL}/xgb-ci.${WHEEL_TAG}:${IMAGE_TAG}
 
 echo "--- Build CPU code targeting ARM64"
 set -x
+
+python3 ops/script/pypi_variants.py --use-cpu-suffix=0 --require-nccl-dep=0
 python3 ops/docker_run.py \
   --image-uri ${IMAGE_URI} \
-  -- ops/pipeline/build-cpu-arm64-impl.sh
+  -- ops/pipeline/build-python-wheels-arm64-impl.sh
 
 echo "--- Audit binary wheel to ensure it's compliant with ${WHEEL_TAG} standard"
 python3 ops/docker_run.py \

ops/pipeline/build-cpu-arm64.sh renamed to ops/pipeline/build-python-wheels-arm64.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+# Build Python wheels, CPU variant (no federated learning)
+
+set -euo pipefail
+
+if [[ -z "${GITHUB_SHA:-}" ]]
+then
+  echo "Make sure to set environment variable GITHUB_SHA"
+  exit 1
+fi
+
+if [[ "$#" -lt 2 ]]
+then
+  echo "Usage: $0 {manylinux2014,manylinux_2_28} {x86_64,aarch64}"
+  exit 1
+fi
+
+manylinux_target="$1"
+arch="$2"
+
+source ops/pipeline/classify-git-branch.sh
+source ops/pipeline/get-docker-registry-details.sh
+source ops/pipeline/get-image-tag.sh
+
+WHEEL_TAG="${manylinux_target}_${arch}"
+IMAGE_REPO="xgb-ci.${WHEEL_TAG}"
+IMAGE_URI="${DOCKER_REGISTRY_URL}/${IMAGE_REPO}:${IMAGE_TAG}"
+PYTHON_BIN="/opt/python/cp310-cp310/bin/python"
+
+echo "--- Build binary wheel for ${WHEEL_TAG} (CPU only)"
+set -x
+
+# Patch to rename pkg to xgboost-cpu
+python3 ops/script/pypi_variants.py --use-cpu-suffix=1 --require-nccl-dep=0
+python3 ops/docker_run.py \
+  --image-uri "${IMAGE_URI}" \
+  -- bash -c \
+  "cd python-package && ${PYTHON_BIN} -m pip wheel --no-deps -v . --wheel-dir dist/"
+
+python3 ops/docker_run.py \
+  --image-uri "${IMAGE_URI}" \
+  -- auditwheel repair --only-plat \
+  --plat ${WHEEL_TAG} python-package/dist/xgboost_cpu-*.whl
+python3 -m wheel tags --python-tag py3 --abi-tag none --platform ${WHEEL_TAG} --remove \
+  wheelhouse/xgboost_cpu-*.whl
+rm -v python-package/dist/xgboost_cpu-*.whl
+mv -v wheelhouse/xgboost_cpu-*.whl python-package/dist/
+
+if ! unzip -l ./python-package/dist/*.whl | grep libgomp > /dev/null; then
+  echo "error: libgomp.so was not vendored in the wheel"
+  exit -1
+fi
+
+# Check size of wheel
+pydistcheck --config python-package/pyproject.toml python-package/dist/*.whl
+
+if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
+then
+  python3 ops/pipeline/manage-artifacts.py upload \
+    --s3-bucket xgboost-nightly-builds \
+    --prefix ${BRANCH_NAME}/${GITHUB_SHA} --make-public \
+    python-package/dist/*.whl
+fi

ops/pipeline/build-python-wheels-cpu.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+# Build Python wheels targeting MacOS (no federated learning)
 
 set -euox pipefail
 

ops/pipeline/build-python-wheels-macos.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+# Build Python wheels targeting manylinux2014 (no GPU, no federated learning)
 
 set -euo pipefail
 
@@ -28,13 +29,11 @@ PYTHON_BIN="/opt/python/cp310-cp310/bin/python"
 echo "--- Build binary wheel for ${WHEEL_TAG}"
 set -x
 
-python3 ops/script/pypi_variants.py --variant=manylinux2014
+python3 ops/script/pypi_variants.py --use-cpu-suffix=0 --require-nccl-dep=0
 python3 ops/docker_run.py \
   --image-uri "${IMAGE_URI}" \
   -- bash -c \
   "cd python-package && ${PYTHON_BIN} -m pip wheel --no-deps -v . --wheel-dir dist/"
-# discard the patch
-python3 ops/script/pypi_variants.py --variant=default
 
 python3 ops/docker_run.py \
   --image-uri "${IMAGE_URI}" \
@@ -46,25 +45,6 @@ rm -rf python-package/dist/
 mkdir python-package/dist/
 mv -v wheelhouse/*.whl python-package/dist/
 
-echo "--- Build binary wheel for ${WHEEL_TAG} (CPU only)"
-# Patch to rename pkg to xgboost-cpu
-python3 ops/script/pypi_variants.py --variant=cpu
-python3 ops/docker_run.py \
-  --image-uri "${IMAGE_URI}" \
-  -- bash -c \
-  "cd python-package && ${PYTHON_BIN} -m pip wheel --no-deps -v . --wheel-dir dist/"
-# discard the patch
-python3 ops/script/pypi_variants.py --variant=default
-
-python3 ops/docker_run.py \
-  --image-uri "${IMAGE_URI}" \
-  -- auditwheel repair --only-plat \
-  --plat ${WHEEL_TAG} python-package/dist/xgboost_cpu-*.whl
-python3 -m wheel tags --python-tag py3 --abi-tag none --platform ${WHEEL_TAG} --remove \
-  wheelhouse/xgboost_cpu-*.whl
-rm -v python-package/dist/xgboost_cpu-*.whl
-mv -v wheelhouse/xgboost_cpu-*.whl python-package/dist/
-
 if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
 then
   python3 ops/pipeline/manage-artifacts.py upload \

ops/pipeline/build-manylinux2014.sh renamed to ops/pipeline/build-python-wheels-manylinux2014.sh

@@ -18,7 +18,7 @@ Write-Host "--- Build binary wheel"
 cd ..
 # Patch to rename pkg to xgboost-cpu
 conda activate
-python ops/script/pypi_variants.py --variant=cpu
+python ops/script/pypi_variants.py --use-cpu-suffix=1 --require-nccl-dep=0
 if ($LASTEXITCODE -ne 0) { throw "Last command failed" }
 
 cd python-package

ops/pipeline/build-win64-cpu.ps1

@@ -61,7 +61,7 @@ def pypkg(
     with open(pyprj_path, "w") as fd:
         fd.write(pyprj)
 
-    make_pyproject("default")
+    make_pyproject(use_cpu_suffix=0, require_nccl_dep=1)
 
 
 @cd(R_PACKAGE)

ops/script/change_version.py

@@ -3,17 +3,15 @@
 import argparse
 import os
 
-from test_utils import PY_PACKAGE, ROOT
+from test_utils import PY_PACKAGE
 
 IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.in")
 OUT_PATH = os.path.join(PY_PACKAGE, "pyproject.toml")
 
-WHL_CPU = """
-[tool.hatch.build.targets.wheel]
-packages = ["xgboost/"]
-"""
+NCCL_WHL = """    \"nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'\","""
 
-CHOICES = ["default", "cpu", "manylinux2014"]
+NAME = "{{ name }}"
+NCCL = "{{ nccl }}"
 
 
 def copyfile(src: str, dst: str) -> None:
@@ -23,39 +21,22 @@ def copyfile(src: str, dst: str) -> None:
         fd.write(content)
 
 
-def make_pyproject(variant: str) -> None:
-    assert variant in CHOICES
+def make_pyproject(*, use_cpu_suffix: int, require_nccl_dep: int) -> None:
+    if use_cpu_suffix == 1 and require_nccl_dep == 1:
+        raise ValueError(
+            "xgboost-cpu cannot require NCCL dependency. "
+            "If --use-cpu-suffix=1, you must set --require-nccl-dep=0."
+        )
 
     with open(IN_PATH) as fd:
         pyproject = fd.read()
 
     readme_dft = os.path.join(PY_PACKAGE, "README.dft.rst")
     readme_cpu = os.path.join(PY_PACKAGE, "README.cpu.rst")
     readme = os.path.join(PY_PACKAGE, "README.rst")
-    if variant == "cpu":
-        pyproject = (
-            pyproject.replace("{{ name }}", "xgboost-cpu")
-            .replace("{{ wheel }}", WHL_CPU)
-            .replace("{{ nccl }}", "")
-        )
-        copyfile(readme_cpu, readme)
-    elif variant == "manylinux2014":
-        pyproject = (
-            pyproject.replace("{{ name }}", "xgboost")
-            .replace("{{ wheel }}", "")
-            .replace("{{ nccl }}", "")
-        )
-        copyfile(readme_dft, readme)
-    else:
-        pyproject = (
-            pyproject.replace("{{ name }}", "xgboost")
-            .replace("{{ wheel }}", "")
-            .replace(
-                "{{ nccl }}",
-                """    \"nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'\",""",
-            )
-        )
-        copyfile(readme_dft, readme)
+    pyproject = pyproject.replace(NAME, "xgboost-cpu" if use_cpu_suffix else "xgboost")
+    copyfile(readme_cpu if use_cpu_suffix else readme_dft, readme)
+    pyproject = pyproject.replace(NCCL, NCCL_WHL if require_nccl_dep else "")
     pyproject = (
         f"# Generated by `{os.path.basename(__file__)}`, don't edit.\n" + pyproject
     )
@@ -67,10 +48,21 @@ def make_pyproject(variant: str) -> None:
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--variant",
-        type=str,
-        choices=CHOICES,
-        default="default",
+        "--use-cpu-suffix",
+        type=int,
+        choices=[0, 1],
+        required=True,
+        help="Whether to rename the package name to xgboost-cpu",
+    )
+    parser.add_argument(
+        "--require-nccl-dep",
+        type=int,
+        choices=[0, 1],
+        required=True,
+        help="Whether to require the NCCL dependency",
     )
     args = parser.parse_args()
-    make_pyproject(args.variant)
+    make_pyproject(
+        use_cpu_suffix=args.use_cpu_suffix,
+        require_nccl_dep=args.require_nccl_dep,
+    )