Set up CI/CD pipeline for CI containers

.github/runs-on.yml

@@ -0,0 +1,29 @@
+# Custom images with CUDA toolkit installed
+# See vm_images/ for instructions for building the images
+images:
+  linux-amd64:
+    platform: "linux"
+    arch: "x64"
+    owner: "492475357299"  # XGBooost CI
+    name: "xgboost-ci-runs-on-linux-*"
+  windows-amd64:
+    platform: "windows"
+    arch: "x64"
+    owner: "492475357299"  # XGBooost CI
+    name: "xgboost-ci-runs-on-windows-*"
+
+runners:
+  linux-amd64-cpu:
+    cpu: 16
+    family: ["c7i-flex", "c7i", "c7a", "c5", "c5a"]
+    image: linux-amd64
+  linux-amd64-gpu:
+    family: ["g4dn.xlarge"]
+    image: linux-amd64
+  linux-amd64-mgpu:
+    family: ["g4dn.12xlarge"]
+    image: linux-amd64
+  linux-arm64-cpu:
+    cpu: 16
+    family: ["c6g", "c7g"]
+    image: ubuntu24-full-arm64

.github/workflows/containers.yml

@@ -0,0 +1,63 @@
+name: Build and publish image
+run-name: "Build image${{ (inputs.upstream_repository != '') && format(' - triggered by: {0}', inputs.upstream_repository) || '' }}"
+
+on:
+  workflow_dispatch:
+    inputs:
+      upstream_repository:
+        required: false
+        type: string
+      upstream_job:
+        required: false
+        type: string
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+    - cron: "0 7 * * *"  # Run once daily
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  BRANCH_NAME: >-
+    ${{ github.event.pull_request.number && 'PR-' }}${{ github.event.pull_request.number || github.ref_name }}
+  PUBLISH_CONTAINER: 1
+
+jobs:
+  build-containers:
+    name: Build CI containers (${{ matrix.container_id }})
+    runs-on:
+      - runs-on
+      - runner=${{ matrix.runner }}
+      - run-id=${{ github.run_id }}
+      - tag=build-containers-${{ matrix.container_id }}
+    strategy:
+      matrix:
+        container_id:
+          - xgb-ci.clang_tidy
+          - xgb-ci.cpu
+          - xgb-ci.gpu
+          - xgb-ci.gpu_build_r_rockylinux8
+          - xgb-ci.gpu_build_rockylinux8
+          - xgb-ci.gpu_build_rockylinux8_dev_ver
+          - xgb-ci.jvm
+          - xgb-ci.jvm_gpu_build
+          - xgb-ci.manylinux_2_28_x86_64
+          - xgb-ci.manylinux2014_x86_64
+        runner: [linux-amd64-cpu]
+        include:
+          - container_id: xgb-ci.aarch64
+            runner: linux-arm64-cpu
+          - container_id: xgb-ci.manylinux2014_aarch64
+            runner: linux-arm64-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: Build ${{ matrix.container_id }}
+        run: bash containers/docker_build.sh ${{ matrix.container_id }}

.gitignore

@@ -0,0 +1 @@
+__pycache__

containers/ci_container.yml

@@ -0,0 +1,72 @@
+## List of CI containers with definitions and build arguments
+
+# Each container will be built using the definition from
+# containers/dockerfile/Dockerfile.CONTAINER_DEF
+
+rapids_versions:
+  stable: &rapids_version "24.10"
+  dev: &dev_rapids_version "24.12"
+
+xgb-ci.gpu_build_rockylinux8:
+  container_def: gpu_build_rockylinux8
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    NCCL_VERSION: "2.23.4-1"
+    RAPIDS_VERSION: *rapids_version
+
+xgb-ci.gpu_build_rockylinux8_dev_ver:
+  container_def: gpu_build_rockylinux8
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    NCCL_VERSION: "2.23.4-1"
+    RAPIDS_VERSION: *dev_rapids_version
+
+xgb-ci.gpu_build_r_rockylinux8:
+  container_def: gpu_build_r_rockylinux8
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    R_VERSION: "4.3.2"
+
+xgb-ci.gpu:
+  container_def: gpu
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    NCCL_VERSION: "2.23.4-1"
+    RAPIDS_VERSION: *rapids_version
+
+xgb-ci.gpu_dev_ver:
+  container_def: gpu
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    NCCL_VERSION: "2.23.4-1"
+    RAPIDS_VERSION: *dev_rapids_version
+    RAPIDSAI_CONDA_CHANNEL: "rapidsai-nightly"
+
+xgb-ci.clang_tidy:
+  container_def: clang_tidy
+  build_args:
+    CUDA_VERSION: "12.4.1"
+
+xgb-ci.cpu:
+  container_def: cpu
+
+xgb-ci.aarch64:
+  container_def: aarch64
+
+xgb-ci.manylinux_2_28_x86_64:
+  container_def: manylinux_2_28_x86_64
+
+xgb-ci.manylinux2014_x86_64:
+  container_def: manylinux2014_x86_64
+
+xgb-ci.manylinux2014_aarch64:
+  container_def: manylinux2014_aarch64
+
+xgb-ci.jvm:
+  container_def: jvm
+
+xgb-ci.jvm_gpu_build:
+  container_def: jvm_gpu_build
+  build_args:
+    CUDA_VERSION: "12.4.1"
+    NCCL_VERSION: "2.23.4-1"

containers/docker_build.py

@@ -0,0 +1,140 @@
+"""
+Wrapper script to build a Docker container
+"""
+
+import argparse
+import itertools
+import pathlib
+import subprocess
+import sys
+import textwrap
+
+PROJECT_ROOT_DIR = pathlib.Path(__file__).parent.parent
+LINEWIDTH = 88
+TEXT_WRAPPER = textwrap.TextWrapper(
+    width=LINEWIDTH,
+    initial_indent="",
+    subsequent_indent="    ",
+    break_long_words=False,
+    break_on_hyphens=False,
+)
+
+
+def fancy_print_cli_args(cli_args: list[str]) -> None:
+    print(
+        "=" * LINEWIDTH
+        + "\n"
+        + "  \\\n".join(TEXT_WRAPPER.wrap(" ".join(cli_args)))
+        + "\n"
+        + "=" * LINEWIDTH
+        + "\n",
+        flush=True,
+    )
+
+
+def parse_build_args(*, raw_build_args: list[str]) -> dict[str, str]:
+    parsed_build_args = dict()
+    for arg in raw_build_args:
+        try:
+            key, value = arg.split("=", maxsplit=1)
+        except ValueError as e:
+            raise ValueError(
+                f"Build argument must be of form KEY=VALUE. Got: {arg}"
+            ) from e
+        parsed_build_args[key] = value
+    return parsed_build_args
+
+
+def docker_build(
+    *,
+    container_tag: str,
+    build_args: dict[str, str],
+    dockerfile_path: pathlib.Path,
+    docker_context_path: pathlib.Path,
+) -> None:
+    ## Set up command-line arguments to be passed to `docker build`
+    # Build args
+    docker_build_cli_args = list(
+        itertools.chain.from_iterable(
+            [["--build-arg", f"{k}={v}"] for k, v in build_args.items()]
+        )
+    )
+    # When building an image using a non-default driver, we need to specify
+    # `--load` to load it to the image store.
+    # See https://docs.docker.com/build/builders/drivers/
+    docker_build_cli_args.append("--load")
+    # Remaining CLI args
+    docker_build_cli_args.extend(
+        [
+            "--progress=plain",
+            "--ulimit",
+            "nofile=1024000:1024000",
+            "-t",
+            container_tag,
+            "-f",
+            str(dockerfile_path),
+            str(docker_context_path),
+        ]
+    )
+    cli_args = ["docker", "build"] + docker_build_cli_args
+    fancy_print_cli_args(cli_args)
+    subprocess.run(cli_args, check=True, encoding="utf-8")
+
+
+def main(*, args: argparse.Namespace) -> None:
+    docker_context_path = PROJECT_ROOT_DIR / "containers"
+    dockerfile_path = (
+        docker_context_path / "dockerfile" / f"Dockerfile.{args.container_def}"
+    )
+
+    build_args = parse_build_args(raw_build_args=args.build_arg)
+
+    docker_build(
+        container_tag=args.container_tag,
+        build_args=build_args,
+        dockerfile_path=dockerfile_path,
+        docker_context_path=docker_context_path,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Build a Docker container")
+    parser.add_argument(
+        "--container-def",
+        type=str,
+        required=True,
+        help=(
+            "String uniquely identifying the container definition. The container "
+            "definition will be fetched from "
+            "containers/dockerfile/Dockerfile.CONTAINER_DEF."
+        ),
+    )
+    parser.add_argument(
+        "--container-tag",
+        type=str,
+        required=True,
+        help=(
+            "Tag to assign to the newly built container, e.g. "
+            "492475357299.dkr.ecr.us-west-2.amazonaws.com/xgb-ci.gpu:master"
+        ),
+    )
+    parser.add_argument(
+        "--build-arg",
+        type=str,
+        default=[],
+        action="append",
+        help=(
+            "Build-time variable(s) to be passed to `docker build`. Each variable "
+            "should be specified as a key-value pair in the form KEY=VALUE. "
+            "The variables should match the ARG instructions in the Dockerfile. "
+            "When passing multiple variables, specify --build-arg multiple times. "
+            "Example: --build-arg CUDA_VERSION_ARG=12.5 --build-arg RAPIDS_VERSION_ARG=24.10"
+        ),
+    )
+
+    if len(sys.argv) == 1:
+        parser.print_help()
+        sys.exit(1)
+
+    parsed_args = parser.parse_args()
+    main(args=parsed_args)