[CI] Build variant wheels for CUDA 12 and 13

.github/workflows/python_wheels_winarm64.yml

@@ -46,7 +46,7 @@ jobs:
       - name: Build XGBoost Python wheel for Win-ARM64
         run: |
           # Patch to rename pkg to xgboost-cpu
-          python ops/script/pypi_variants.py --use-cpu-suffix=1 --require-nccl-dep=0
+          python ops/script/pypi_variants.py --use-suffix=cpu --require-nccl-dep=na
           cd python-package
           mkdir -p wheelhouse
           pip wheel --no-deps -v . --wheel-dir wheelhouse/

ops/pipeline/build-cuda-impl.sh

@@ -11,14 +11,14 @@ else
   cmake_args=''
 fi
 
-if [[ "${USE_FEDERATED:-}" == 1 ]]
+if [[ "${USE_FEDERATED:-0}" == 1 ]]
 then
   cmake_args="${cmake_args} -DPLUGIN_FEDERATED=ON"
 else
   cmake_args="${cmake_args} -DPLUGIN_FEDERATED=OFF"
 fi
 
-if [[ "${USE_RMM:-}" == 1 ]]
+if [[ "${USE_RMM:-0}" == 1 ]]
 then
   cmake_prefix_path='/opt/grpc;/opt/rmm;/opt/rmm/lib64/rapids/cmake'
   cmake_args="${cmake_args} -DPLUGIN_RMM=ON"

ops/pipeline/build-cuda13.sh

@@ -32,9 +32,7 @@ fi
 
 set -x
 
-# Remove nvidia-nccl-cu12 from the list of Python deps
-# nvidia-nccl-cu13 is not yet available on PyPI
-python3 ops/script/pypi_variants.py --use-cpu-suffix=0 --require-nccl-dep=0
+python3 ops/script/pypi_variants.py --use-suffix=cu13 --require-nccl-dep=cu13
 
 python3 ops/docker_run.py \
   --image-uri ${BUILD_IMAGE_URI} \
@@ -56,3 +54,12 @@ fi
 
 # Check size of wheel
 pydistcheck --config python-package/pyproject.toml python-package/dist/*.whl
+
+echo "--- Upload Python wheel"
+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-arm64.sh

@@ -19,7 +19,7 @@ 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/script/pypi_variants.py --use-suffix=na --require-nccl-dep=na
 python3 ops/docker_run.py \
   --image-uri ${IMAGE_URI} \
   -- ops/pipeline/build-python-wheels-arm64-impl.sh

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

@@ -31,7 +31,7 @@ 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/script/pypi_variants.py --use-suffix=cpu --require-nccl-dep=na
 python3 ops/docker_run.py \
   --image-uri "${IMAGE_URI}" \
   -- bash -c \

ops/pipeline/build-variant-wheels.sh

@@ -10,6 +10,8 @@ then
 fi
 
 image_repo='xgb-ci.gpu_build_rockylinux8'
+export USE_RMM=0
+export USE_FEDERATED=0
 
 source ops/pipeline/classify-git-branch.sh
 source ops/pipeline/get-docker-registry-details.sh
@@ -27,13 +29,12 @@ then
 else
   export BUILD_ONLY_SM75=0
 fi
-export USE_RMM=0
 
 set -x
 
 python3 ops/docker_run.py \
   --image-uri ${BUILD_IMAGE_URI} \
-  --run-args='-e BUILD_ONLY_SM75 -e USE_RMM' \
+  --run-args='-e BUILD_ONLY_SM75 -e USE_RMM -e USE_FEDERATED' \
   -- ops/pipeline/build-cuda-impl.sh
 
 echo "--- Audit binary wheel to ensure it's compliant with ${WHEEL_TAG} standard"

ops/pipeline/build-win64-cpu.ps1

@@ -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 --use-cpu-suffix=1 --require-nccl-dep=0
+python ops/script/pypi_variants.py --use-suffix=cpu --require-nccl-dep=na
 if ($LASTEXITCODE -ne 0) { throw "Last command failed" }
 
 cd python-package

ops/script/change_version.py

@@ -52,16 +52,17 @@ def pypkg(
     with open(pyver_path, "w") as fd:
         fd.write(pyver + "\n")
 
-    pyprj_path = os.path.join("pyproject.toml.in")
-    with open(pyprj_path, "r") as fd:
-        pyprj = fd.read()
-    matched = re.search('version = "' + r"([0-9]+\.[0-9]+\.[0-9]+.*)" + '"', pyprj)
-    assert matched, "Couldn't find version string in pyproject.toml."
-    pyprj = pyprj[: matched.start(1)] + pyver + pyprj[matched.end(1) :]
-    with open(pyprj_path, "w") as fd:
-        fd.write(pyprj)
-
-    make_pyproject(use_cpu_suffix=0, require_nccl_dep=1)
+    for pyprj_file in ["pyproject.toml.in", "pyproject.toml.stub.in"]:
+        pyprj_path = os.path.join(pyprj_file)
+        with open(pyprj_path, "r") as fd:
+            pyprj = fd.read()
+        matched = re.search('version = "' + r"([0-9]+\.[0-9]+\.[0-9]+.*)" + '"', pyprj)
+        assert matched, "Couldn't find version string in pyproject.toml."
+        pyprj = pyprj[: matched.start(1)] + pyver + pyprj[matched.end(1) :]
+        with open(pyprj_path, "w") as fd:
+            fd.write(pyprj)
+
+    make_pyproject(use_suffix="na", require_nccl_dep="cu12")
 
 
 @cd(R_PACKAGE)

ops/script/pypi_variants.py

@@ -2,16 +2,21 @@
 
 import argparse
 import os
+import tomllib
 
+from packaging.version import Version
 from test_utils import PY_PACKAGE
 
 IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.in")
+STUB_IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.stub.in")
 OUT_PATH = os.path.join(PY_PACKAGE, "pyproject.toml")
 
-NCCL_WHL = """    \"nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'\","""
+NCCL_WHL = """    \"nvidia-nccl-{0} ; platform_system == 'Linux' and platform_machine != 'aarch64'\","""
 
 NAME = "{{ name }}"
 NCCL = "{{ nccl }}"
+VERSION = "{{ version }}"
+CUDA_VARIANTS = ["cu12", "cu13"]
 
 
 def copyfile(src: str, dst: str) -> None:
@@ -21,22 +26,55 @@ def copyfile(src: str, dst: str) -> None:
         fd.write(content)
 
 
-def make_pyproject(*, use_cpu_suffix: int, require_nccl_dep: int) -> None:
-    if use_cpu_suffix == 1 and require_nccl_dep == 1:
+def make_pyproject(
+    *, use_suffix: str, require_nccl_dep: str, create_stub: bool = False
+) -> None:
+    if use_suffix == "cpu" and require_nccl_dep != "na":
         raise ValueError(
             "xgboost-cpu cannot require NCCL dependency. "
-            "If --use-cpu-suffix=1, you must set --require-nccl-dep=0."
+            "When setting --use-suffix='cpu', you must also set --require-nccl-dep='na'."
         )
+    if (
+        use_suffix in CUDA_VARIANTS
+        and require_nccl_dep in CUDA_VARIANTS
+        and use_suffix != require_nccl_dep
+    ):
+        raise ValueError(
+            "Inconsistent choices for --use-suffix and --require-nccl-dep. "
+            "When --use-suffix is set to one of {{{0}}}, --require-nccl-dep must be "
+            "set to identical value as --use-suffix.".format(",".join(CUDA_VARIANTS))
+        )
+    if create_stub:
+        if use_suffix == "na":
+            raise ValueError("To create a stub package, --use-suffix must not be 'na'")
+        if require_nccl_dep != "na":
+            raise ValueError(
+                "To create a stub package, --require-nccl-dep must be 'na'"
+            )
 
-    with open(IN_PATH) as fd:
+    with open(STUB_IN_PATH if create_stub else 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_stub = os.path.join(PY_PACKAGE, "README.stub.rst")
     readme = os.path.join(PY_PACKAGE, "README.rst")
-    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 = pyproject.replace(
+        NAME, f"xgboost-{use_suffix}" if use_suffix != "na" else "xgboost"
+    )
+    if create_stub:
+        copyfile(readme_stub, readme)
+        pyproject_parsed = tomllib.loads(pyproject)
+        pyproject = pyproject.replace(
+            VERSION, str(Version(pyproject_parsed["project"]["version"]))
+        )
+    elif use_suffix == "cpu":
+        copyfile(readme_cpu, readme)
+    else:
+        copyfile(readme_dft, readme)
+    pyproject = pyproject.replace(
+        NCCL, NCCL_WHL.format(require_nccl_dep) if require_nccl_dep != "na" else ""
+    )
     pyproject = (
         f"# Generated by `{os.path.basename(__file__)}`, don't edit.\n" + pyproject
     )
@@ -48,21 +86,27 @@ def make_pyproject(*, use_cpu_suffix: int, require_nccl_dep: int) -> None:
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--use-cpu-suffix",
-        type=int,
-        choices=[0, 1],
-        required=True,
-        help="Whether to rename the package name to xgboost-cpu",
+        "--use-suffix",
+        type=str,
+        choices=["na", "cpu"] + CUDA_VARIANTS,
+        default="na",
+        help="When using this option, rename the package name to xgboost-[suffix]. Set to 'na' to disable",
     )
     parser.add_argument(
         "--require-nccl-dep",
-        type=int,
-        choices=[0, 1],
+        type=str,
+        choices=["na"] + CUDA_VARIANTS,
         required=True,
-        help="Whether to require the NCCL dependency",
+        help="Which NCCL dependency to use; select 'na' to remove NCCL dependency",
+    )
+    parser.add_argument(
+        "--create-stub",
+        action="store_true",
+        help="Create a stub package that redirects users to install `xgboost`",
     )
     args = parser.parse_args()
     make_pyproject(
-        use_cpu_suffix=args.use_cpu_suffix,
+        use_suffix=args.use_suffix,
         require_nccl_dep=args.require_nccl_dep,
+        create_stub=args.create_stub,
     )

ops/script/release_artifacts.py

@@ -111,7 +111,7 @@ def make_python_sdist(
 
     # Build sdist for `xgboost-cpu`.
     with DirectoryExcursion(ROOT):
-        make_pyproject(use_cpu_suffix=1, require_nccl_dep=0)
+        make_pyproject(use_suffix="cpu", require_nccl_dep="na")
     with DirectoryExcursion(ROOT / "python-package"):
         subprocess.run(["python", "-m", "build", "--sdist"], check=True)
         sdist_name = (
@@ -126,7 +126,7 @@ def make_python_sdist(
 
     # Build sdist for `xgboost`.
     with DirectoryExcursion(ROOT):
-        make_pyproject(use_cpu_suffix=0, require_nccl_dep=1)
+        make_pyproject(use_suffix="na", require_nccl_dep="cu12")
 
     with DirectoryExcursion(ROOT / "python-package"):
         subprocess.run(["python", "-m", "build", "--sdist"], check=True)
@@ -140,6 +140,22 @@ def make_python_sdist(
         dest = dist_dir / sdist_name
         shutil.move(src, dest)
 
+    # Build stub package `xgboost-cu12`.
+    with DirectoryExcursion(ROOT):
+        make_pyproject(use_suffix="cu12", require_nccl_dep="na", create_stub=True)
+
+    with DirectoryExcursion(ROOT / "python-package"):
+        subprocess.run(["python", "-m", "build", "--sdist"], check=True)
+        sdist_name = (
+            f"xgboost_cu12-{release}{rc}{rc_ver}.tar.gz"
+            if rc
+            else f"xgboost_cu12-{release}.tar.gz"
+        )
+        src = DIST / sdist_name
+        subprocess.run(["twine", "check", str(src)], check=True)
+        dest = dist_dir / sdist_name
+        shutil.move(src, dest)
+
 
 def download_python_wheels(branch: str, commit_hash: str, outdir: Path) -> None:
     """Download all Python binary wheels for the specified branch."""

python-package/README.stub.rst

@@ -0,0 +1,5 @@
+======================================
+Placeholder for XGBoost Python Package
+======================================
+
+This package is a placeholder for the `xgboost` package.

python-package/pyproject.toml.stub.in

@@ -0,0 +1,38 @@
+[build-system]
+requires = [
+    "hatchling>=1.12.1",
+]
+build-backend = "hatchling.build"
+
+[project]
+name = "{{ name }}"
+description = "XGBoost Python Package"
+readme = { file = "README.rst", content-type = "text/x-rst" }
+authors = [
+    { name = "Hyunsu Cho", email = "[email protected]" },
+    { name = "Jiaming Yuan", email = "[email protected]" }
+]
+version = "3.1.0-dev"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Development Status :: 5 - Production/Stable",
+    "Operating System :: OS Independent",
+    "Typing :: Typed",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+    "xgboost=={{ version }}",
+]
+
+[tool.hatch.build.targets.sdist]
+only-include = ["pyproject.toml"]
+
+[tool.hatch.build.targets.wheel]
+only-include = ["pyproject.toml"]