[CI] Build a CPU-only wheel under name xgboost-cpu (#10603)

dev/release-artifacts.py

@@ -2,6 +2,7 @@
 
 tqdm, sh are required to run this script.
 """
+
 import argparse
 import os
 import shutil
@@ -106,6 +107,15 @@ def make_pysrc_wheel(
     if not os.path.exists(dist):
         os.mkdir(dist)
 
+    # Apply patch to remove NCCL dependency
+    # Save the original content of pyproject.toml so that we can restore it later
+    with DirectoryExcursion(ROOT):
+        with open("python-package/pyproject.toml", "r") as f:
+            orig_pyproj_lines = f.read()
+        with open("tests/buildkite/remove_nccl_dep.patch", "r") as f:
+            patch_lines = f.read()
+        subprocess.run(["patch", "-p0"], input=patch_lines, text=True)
+
     with DirectoryExcursion(os.path.join(ROOT, "python-package")):
         subprocess.check_call(["python", "-m", "build", "--sdist"])
         if rc is not None:
@@ -117,6 +127,10 @@ def make_pysrc_wheel(
         target = os.path.join(dist, name)
         shutil.move(src, target)
 
+    with DirectoryExcursion(ROOT):
+        with open("python-package/pyproject.toml", "w") as f:
+            print(orig_pyproj_lines, file=f, end="")
+
 
 def download_py_packages(
     branch: str, major: int, minor: int, commit_hash: str, outdir: str

doc/install.rst

@@ -76,6 +76,19 @@ Capabilities of binary wheels for each platform:
 | Windows             | |tick|  |  |cross|             |
 +---------------------+---------+----------------------+
 
+Minimal installation (CPU-only)
+*******************************
+The default installation with ``pip`` will install the full XGBoost package, including the support for the GPU algorithms and federated learning.
+
+You may choose to reduce the size of the installed package and save the disk space, by opting to install ``xgboost-cpu`` instead:
+
+.. code-block:: bash
+
+  pip install xgboost-cpu
+
+The ``xgboost-cpu`` variant will have drastically smaller disk footprint, but does not provide some features, such as the GPU algorithms and
+federated learning.
+
 Conda
 *****
 

python-package/pyproject.toml

@@ -7,13 +7,13 @@ build-backend = "packager.pep517"
 
 [project]
 name = "xgboost"
-version = "2.2.0-dev"
+description = "XGBoost Python Package"
+readme = { file = "README.rst", content-type = "text/x-rst" }
 authors = [
     { name = "Hyunsu Cho", email = "chohyu01@cs.washington.edu" },
     { name = "Jiaming Yuan", email = "jm.yuan@outlook.com" }
 ]
-description = "XGBoost Python Package"
+version = "2.2.0-dev"
-readme = { file = "README.rst", content-type = "text/x-rst" }
 requires-python = ">=3.8"
 license = { text = "Apache-2.0" }
 classifiers = [

tests/buildkite/build-manylinux2014-aarch64.sh

@@ -1,33 +0,0 @@
-#!/bin/bash
-
-set -euo pipefail
-
-source tests/buildkite/conftest.sh
-
-WHEEL_TAG=manylinux2014_aarch64
-command_wrapper="tests/ci_build/ci_build.sh manylinux2014_aarch64"
-python_bin="/opt/python/cp310-cp310/bin/python"
-
-echo "--- Build binary wheel for ${WHEEL_TAG}"
-# Patch to add warning about manylinux2014 variant
-patch -p0 < tests/buildkite/manylinux2014_warning.patch
-$command_wrapper bash -c \
-  "cd python-package && ${python_bin} -m pip wheel --no-deps -vvv . --wheel-dir dist/"
-git checkout python-package/xgboost/core.py  # discard the patch
-
-$command_wrapper auditwheel repair --plat ${WHEEL_TAG} python-package/dist/*.whl
-$command_wrapper ${python_bin} tests/ci_build/rename_whl.py  \
-  --wheel-path wheelhouse/*.whl  \
-  --commit-hash ${BUILDKITE_COMMIT}  \
-  --platform-tag ${WHEEL_TAG}
-rm -rf python-package/dist/
-mkdir python-package/dist/
-mv -v wheelhouse/*.whl python-package/dist/
-
-echo "--- Upload Python wheel"
-buildkite-agent artifact upload python-package/dist/*.whl
-if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
-then
-  aws s3 cp python-package/dist/*.whl s3://xgboost-nightly-builds/${BRANCH_NAME}/ \
-    --acl public-read --no-progress
-fi

tests/buildkite/build-manylinux2014-x86_64.sh

@@ -1,33 +0,0 @@
-#!/bin/bash
-
-set -euo pipefail
-
-source tests/buildkite/conftest.sh
-
-WHEEL_TAG=manylinux2014_x86_64
-command_wrapper="tests/ci_build/ci_build.sh manylinux2014_x86_64"
-python_bin="/opt/python/cp310-cp310/bin/python"
-
-echo "--- Build binary wheel for ${WHEEL_TAG}"
-# Patch to add warning about manylinux2014 variant
-patch -p0 < tests/buildkite/manylinux2014_warning.patch
-$command_wrapper bash -c \
-  "cd python-package && ${python_bin} -m pip wheel --no-deps -vvv . --wheel-dir dist/"
-git checkout python-package/xgboost/core.py  # discard the patch
-
-$command_wrapper auditwheel repair --plat ${WHEEL_TAG} python-package/dist/*.whl
-$command_wrapper ${python_bin} tests/ci_build/rename_whl.py  \
-  --wheel-path wheelhouse/*.whl  \
-  --commit-hash ${BUILDKITE_COMMIT}  \
-  --platform-tag ${WHEEL_TAG}
-rm -rf python-package/dist/
-mkdir python-package/dist/
-mv -v wheelhouse/*.whl python-package/dist/
-
-echo "--- Upload Python wheel"
-buildkite-agent artifact upload python-package/dist/*.whl
-if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
-then
-  aws s3 cp python-package/dist/*.whl s3://xgboost-nightly-builds/${BRANCH_NAME}/ \
-    --acl public-read --no-progress
-fi

tests/buildkite/build-manylinux2014.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+
+set -euo pipefail
+
+if [ $# -ne 1 ]; then
+  echo "Usage: $0 {x86_64,aarch64}"
+  exit 1
+fi
+
+arch=$1
+
+source tests/buildkite/conftest.sh
+
+WHEEL_TAG="manylinux2014_${arch}"
+command_wrapper="tests/ci_build/ci_build.sh ${WHEEL_TAG}"
+python_bin="/opt/python/cp310-cp310/bin/python"
+
+echo "--- Build binary wheel for ${WHEEL_TAG}"
+# Patch to add warning about manylinux2014 variant
+patch -p0 < tests/buildkite/remove_nccl_dep.patch
+patch -p0 < tests/buildkite/manylinux2014_warning.patch
+$command_wrapper bash -c \
+  "cd python-package && ${python_bin} -m pip wheel --no-deps -v . --wheel-dir dist/"
+git checkout python-package/pyproject.toml python-package/xgboost/core.py  # discard the patch
+
+$command_wrapper auditwheel repair --plat ${WHEEL_TAG} python-package/dist/*.whl
+$command_wrapper ${python_bin} tests/ci_build/rename_whl.py  \
+  --wheel-path wheelhouse/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
+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
+patch -p0 < tests/buildkite/remove_nccl_dep.patch
+patch -p0 < tests/buildkite/cpu_only_pypkg.patch
+$command_wrapper bash -c \
+  "cd python-package && ${python_bin} -m pip wheel --no-deps -v . --wheel-dir dist/"
+git checkout python-package/pyproject.toml  # discard the patch
+
+$command_wrapper auditwheel repair --plat ${WHEEL_TAG} python-package/dist/xgboost_cpu-*.whl
+$command_wrapper ${python_bin} tests/ci_build/rename_whl.py  \
+  --wheel-path wheelhouse/xgboost_cpu-*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
+rm -v python-package/dist/xgboost_cpu-*.whl
+mv -v wheelhouse/xgboost_cpu-*.whl python-package/dist/
+
+echo "--- Upload Python wheel"
+for wheel in python-package/dist/*.whl
+do
+  buildkite-agent artifact upload "${wheel}"
+done
+if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
+then
+  for wheel in python-package/dist/*.whl
+  do
+    aws s3 cp "${wheel}" s3://xgboost-nightly-builds/${BRANCH_NAME}/ \
+      --acl public-read --no-progress
+  done
+fi

tests/buildkite/cpu_only_pypkg.patch

@@ -0,0 +1,55 @@
+diff --git python-package/README.rst python-package/README.rst
+index 1fc0bb5a0..f1c68470b 100644
+--- python-package/README.rst
++++ python-package/README.rst
+@@ -1,20 +1,15 @@
+-======================
+-XGBoost Python Package
+-======================
++=================================
++XGBoost Python Package (CPU only)
++=================================
+ 
+ |PyPI version|
+ 
+-Installation
+-============
++The ``xgboost-cpu`` package provides for a minimal installation, with no support for the GPU algorithms
++or federated learning. It is provided to allow XGBoost to be installed in a space-constrained
++environments.
+ 
+-From `PyPI <https://pypi.python.org/pypi/xgboost>`_
+----------------------------------------------------
++Note. ``xgboost-cpu`` package is only provided for x86_64 (amd64) Linux and Windows platforms.
++For other platforms, please install ``xgboost`` from https://pypi.org/project/xgboost/.
+ 
+-For a stable version, install using ``pip``::
+-
+-    pip install xgboost
+-
+-.. |PyPI version| image:: https://badge.fury.io/py/xgboost.svg
+-   :target: http://badge.fury.io/py/xgboost
+-
+-For building from source, see `build <https://xgboost.readthedocs.io/en/latest/build.html>`_.
++Note. ``xgboost-cpu`` does not provide an sdist (source distribution). You may install sdist
++from https://pypi.org/project/xgboost/.
+diff --git python-package/pyproject.toml python-package/pyproject.toml
+index 46c1451c2..c5dc908d9 100644
+--- python-package/pyproject.toml
++++ python-package/pyproject.toml
+@@ -6,7 +6,7 @@ backend-path = ["."]
+ build-backend = "packager.pep517"
+ 
+ [project]
+-name = "xgboost"
++name = "xgboost-cpu"
+ description = "XGBoost Python Package"
+ readme = { file = "README.rst", content-type = "text/x-rst" }
+ authors = [
+@@ -82,3 +82,6 @@ class-attribute-naming-style = "snake_case"
+ 
+ # Allow single-letter variables
+ variable-rgx = "[a-zA-Z_][a-z0-9_]{0,30}$"
++
++[tool.hatch.build.targets.wheel]
++packages = ["xgboost/"]

tests/buildkite/manylinux2014_warning.patch

@@ -1,17 +1,3 @@
-diff --git python-package/pyproject.toml python-package/pyproject.toml
-index a273d8c13..dee49686a 100644
---- python-package/pyproject.toml
-+++ python-package/pyproject.toml
-@@ -30,8 +30,7 @@ classifiers = [
- ]
- dependencies = [
-     "numpy",
--    "scipy",
--    "nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'"
-+    "scipy"
- ]
- 
- [project.urls]
 diff --git python-package/xgboost/core.py python-package/xgboost/core.py
 index e8bc735e6..030972ef2 100644
 --- python-package/xgboost/core.py

tests/buildkite/pipeline.yml

@@ -72,12 +72,12 @@ steps:
     agents:
       queue: linux-amd64-cpu
   - label: ":console: Build manylinux2014_x86_64 wheel"
-    command: "tests/buildkite/build-manylinux2014-x86_64.sh"
+    command: "tests/buildkite/build-manylinux2014.sh x86_64"
     key: build-manylinux2014-x86_64
     agents:
       queue: linux-amd64-cpu
   - label: ":console: Build manylinux2014_aarch64 wheel"
-    command: "tests/buildkite/build-manylinux2014-aarch64.sh"
+    command: "tests/buildkite/build-manylinux2014.sh aarch64"
     key: build-manylinux2014-aarch64
     agents:
       queue: linux-arm64-cpu

tests/buildkite/remove_nccl_dep.patch

@@ -0,0 +1,14 @@
+diff --git python-package/pyproject.toml python-package/pyproject.toml
+index 8835def25..46c1451c2 100644
+--- python-package/pyproject.toml
++++ python-package/pyproject.toml
+@@ -30,8 +30,7 @@ classifiers = [
+ ]
+ dependencies = [
+     "numpy",
+-    "scipy",
+-    "nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'"
++    "scipy"
+ ]
+ 
+ [project.urls]