From cb1b7a5d168318f4b7a3c1ef40c1dce1f57fe1e1 Mon Sep 17 00:00:00 2001 From: Alexey Rivkin Date: Wed, 8 Jul 2026 13:53:03 +0300 Subject: [PATCH] BUILD: add RAPIDS wrapper-image build tooling Add a declarative compose + one-command buildx-bake script to rebuild the UCXX rapidsai CI wrapper images multi-arch (amd64+arm64) natively, plus a doc for the RAPIDS version-bump procedure. The compose is the single source for the RAPIDS base tags - the script derives its preflight from it. Node endpoints and harbor credentials are supplied via the environment; nothing environment-specific is committed. --- buildlib/dockers/docker-compose-rapidsai.yml | 40 ++++++ buildlib/dockers/rapidsai-images-build.sh | 106 +++++++++++++++ buildlib/dockers/rapidsai-images.md | 131 +++++++++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 buildlib/dockers/docker-compose-rapidsai.yml create mode 100755 buildlib/dockers/rapidsai-images-build.sh create mode 100644 buildlib/dockers/rapidsai-images.md diff --git a/buildlib/dockers/docker-compose-rapidsai.yml b/buildlib/dockers/docker-compose-rapidsai.yml new file mode 100644 index 00000000000..6596176bf60 --- /dev/null +++ b/buildlib/dockers/docker-compose-rapidsai.yml @@ -0,0 +1,40 @@ +# Multi-arch (amd64 + arm64) build of the UCXX rapidsai CI wrapper images. +# +# Driven by rapidsai-images-build.sh (sets up a 2-node buildx builder, then +# `docker buildx bake --push` against this file -> native multi-arch manifests). +# Manual equivalent: +# RAPIDS_VER=26.10 REGISTRY=harbor.mellanox.com docker buildx bake \ +# --builder -f docker-compose-rapidsai.yml \ +# --set '*.platform=linux/amd64,linux/arm64' --push +# +# The platforms come from that --set, not from this file: bake ignores the +# compose `build.platforms` key, so listing it here would be inert. +# +# Bumping RAPIDS: set RAPIDS_VER. If RAPIDS moves the wheel CUDA/py tag +# (e.g. cuda13.3.0 -> 13.3.0, or py3.11 -> py3.12), update the two wheel +# BASE_IMAGE args below AND the matching rapids_cuda_version/py in +# pr/main.yml + azure-pipelines-release.yml. +services: + rapidsai-ci-conda: + image: ${REGISTRY:-harbor.mellanox.com}/ucx/rapidsai-ci-conda:${RAPIDS_VER}-azp-1 + build: + context: . + dockerfile: rapidsai-ci-conda.Dockerfile + args: + BASE_IMAGE: rapidsai/ci-conda:${RAPIDS_VER}-latest + + rapidsai-ci-wheel-cuda13: + image: ${REGISTRY:-harbor.mellanox.com}/ucx/rapidsai-ci-wheel:${RAPIDS_VER}-cuda13-azp-1 + build: + context: . + dockerfile: rapidsai-ci-wheel.Dockerfile + args: + BASE_IMAGE: rapidsai/ci-wheel:${RAPIDS_VER}-cuda13.3.0-rockylinux8-py3.11 + + rapidsai-ci-wheel-cuda12: + image: ${REGISTRY:-harbor.mellanox.com}/ucx/rapidsai-ci-wheel:${RAPIDS_VER}-cuda12-azp-1 + build: + context: . + dockerfile: rapidsai-ci-wheel.Dockerfile + args: + BASE_IMAGE: rapidsai/ci-wheel:${RAPIDS_VER}-cuda12.9.2-rockylinux8-py3.11 diff --git a/buildlib/dockers/rapidsai-images-build.sh b/buildlib/dockers/rapidsai-images-build.sh new file mode 100755 index 00000000000..4a348439e76 --- /dev/null +++ b/buildlib/dockers/rapidsai-images-build.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# +# Build + push the UCXX rapidsai CI wrapper images, multi-arch (amd64 + arm64), +# natively via a 2-node buildx builder + `docker buildx bake` -- one command +# builds both arches on their native node and writes the manifest. +# +# Run on either builder host; it takes the local arch as-is and appends the +# opposite arch over an ssh docker context. Set the ssh endpoint of the OTHER +# builder and the harbor push credentials in the environment: +# +# RAPIDS_VER= \ +# ARM_NODE=ssh://@ # required when running on an x86 host +# X86_NODE=ssh://@ # required when running on an arm host +# HARBOR_UCX_USER=... HARBOR_UCX_PASSWORD=... \ +# ./rapidsai-images-build.sh +# +# The running host must key-auth (non-interactive) ssh to the other node; the +# remote buildkit runs over that ssh context. Do not leave credentials at rest +# on the builder -- pass them via the environment; the script logs out on exit. +# +# Env: +# RAPIDS_VER required, RAPIDS CalVer, e.g. 26.08 +# ARM_NODE ssh endpoint of the arm64 builder (required on an x86 host) +# X86_NODE ssh endpoint of the x86_64 builder (required on an arm host) +# REGISTRY push registry (default: harbor.mellanox.com); also consumed by +# the compose file, so login, push and verify all follow it +# BUILDER buildx builder name (default: ucx-rapidsai) +# HARBOR_UCX_USER / HARBOR_UCX_PASSWORD harbor push credentials (required) + +# Set here rather than on the shebang: the flags are lost when the script is +# invoked as `bash rapidsai-images-build.sh`. +set -eE -o pipefail + +basedir=$(cd "$(dirname "$0")" && pwd) +: "${RAPIDS_VER:?set RAPIDS_VER, e.g. 26.08}" +: "${HARBOR_UCX_USER:?set HARBOR_UCX_USER}" "${HARBOR_UCX_PASSWORD:?set HARBOR_UCX_PASSWORD}" +BUILDER=${BUILDER:-ucx-rapidsai} +REMOTE_CTX=${REMOTE_CTX:-ucx-remote-buildnode} +REGISTRY=${REGISTRY:-harbor.mellanox.com} +export REGISTRY # the compose image: fields resolve it too, so all three agree + +# Run from EITHER builder: the local arch is native, the opposite arch is +# appended over ssh. Only the opposite-arch endpoint is needed. +case "$(uname -m)" in + x86_64) LOCAL_PLAT=linux/amd64; REMOTE_PLAT=linux/arm64 + REMOTE_NODE=${ARM_NODE:?set ARM_NODE=ssh://@} ;; + aarch64) LOCAL_PLAT=linux/arm64; REMOTE_PLAT=linux/amd64 + REMOTE_NODE=${X86_NODE:?set X86_NODE=ssh://@} ;; + *) echo "ERROR: unexpected host arch $(uname -m)" >&2; exit 1 ;; +esac + +# 1. preflight - every BASE_IMAGE in the compose must exist. The compose is +# the single source for the base tags; this parses them instead of keeping a +# second copy that drifts when RAPIDS moves a cuda/py version. +while read -r t; do + t="${t//\$\{RAPIDS_VER\}/$RAPIDS_VER}" + echo ">> preflight: $t" + docker manifest inspect "$t" >/dev/null \ + || { echo "ERROR: base image not found: $t (did RAPIDS move the cuda/py tag?)" >&2; exit 1; } +done < <(sed -n 's/^[[:space:]]*BASE_IMAGE:[[:space:]]*//p' "$basedir/docker-compose-rapidsai.yml") + +# 2. native 2-node buildx builder (idempotent): local node = host arch, remote +# node = opposite arch over ssh. --platform pins each node to ONE arch so buildx +# can't route the foreign arch to the local node's QEMU (crawls on multi-GB images). +# The buildkit image is pinned: newer runc in the floating buildx-stable tag +# fails to start containers on the builders' 5.4-era kernels +# ("can't mask dir /proc/acpi ... invalid argument"). +BUILDKIT_IMAGE=${BUILDKIT_IMAGE:-moby/buildkit:v0.23.2} +if ! docker buildx inspect "$BUILDER" >/dev/null 2>&1; then + docker context inspect "$REMOTE_CTX" >/dev/null 2>&1 \ + || docker context create "$REMOTE_CTX" --docker "host=$REMOTE_NODE" + docker buildx create --name "$BUILDER" --driver docker-container \ + --driver-opt "image=$BUILDKIT_IMAGE" \ + --platform "$LOCAL_PLAT" --bootstrap # local (host arch) + docker buildx create --name "$BUILDER" --append --node "${BUILDER}-remote" \ + --driver-opt "image=$BUILDKIT_IMAGE" \ + --platform "$REMOTE_PLAT" "$REMOTE_CTX" # remote (opposite arch) +fi +docker buildx inspect "$BUILDER" --bootstrap >/dev/null + +# 3. login -> bake (builds both arches on their native node, pushes, writes manifest) -> logout. +# Platforms are forced on the CLI: bake does not honor the compose +# `build.platforms` key and would otherwise build only the local arch. +printf '%s' "$HARBOR_UCX_PASSWORD" | docker login "$REGISTRY" -u "$HARBOR_UCX_USER" --password-stdin +trap 'docker logout "$REGISTRY" >/dev/null 2>&1 || true' EXIT +RAPIDS_VER="$RAPIDS_VER" docker buildx bake \ + --builder "$BUILDER" -f "$basedir/docker-compose-rapidsai.yml" \ + --set '*.platform=linux/amd64,linux/arm64' --push + +# 4. verify the 3 manifests really carry both arches - a single-arch manifest is +# the failure this script exists to prevent, so it has to be fatal, not printed. +echo "=== pushed manifests ===" +for img in rapidsai-ci-conda:${RAPIDS_VER}-azp-1 \ + rapidsai-ci-wheel:${RAPIDS_VER}-cuda12-azp-1 \ + rapidsai-ci-wheel:${RAPIDS_VER}-cuda13-azp-1; do + printf '%-42s ' "$img" + arches=$(docker manifest inspect "$REGISTRY/ucx/$img" \ + | grep -o '"architecture": "[a-z0-9]*"' | grep -o '[a-z0-9]*"$' | tr -d '"' | sort -u) + echo "$arches" | tr '\n' ' ' + for want in amd64 arm64; do + grep -qx "$want" <<< "$arches" \ + || { echo; echo "ERROR: $img is missing $want (manifest is not multi-arch)" >&2; exit 1; } + done + echo +done +echo "DONE: rapidsai-ci-{conda,wheel cuda12/13}:${RAPIDS_VER}-azp-1 (amd64+arm64)" diff --git a/buildlib/dockers/rapidsai-images.md b/buildlib/dockers/rapidsai-images.md new file mode 100644 index 00000000000..c9c5cf13730 --- /dev/null +++ b/buildlib/dockers/rapidsai-images.md @@ -0,0 +1,131 @@ +# UCXX CI: RAPIDS wrapper images + +The UCXX jobs in the openucx/ucx Azure pipelines run inside thin wrapper images +built on top of upstream RAPIDS CI images. This directory holds the compose file +and build script that (re)build and push them. + +The two Dockerfiles they reference, `rapidsai-ci-conda.Dockerfile` and +`rapidsai-ci-wheel.Dockerfile`, arrive with the UCXX CI integration; this tooling +builds nothing until they land. + +| Wrapper image (`ucx/`) | RAPIDS base | Used for | +|--------------------------------------------|--------------------------------------------------------|----------------------------| +| `rapidsai-ci-conda:-azp-1` | `rapidsai/ci-conda:-latest` | conda build + tests | +| `rapidsai-ci-wheel:-cuda13-azp-1` | `rapidsai/ci-wheel:-cuda-rockylinux8-py3.11`| libucxx/ucxx wheels (CUDA 13) | +| `rapidsai-ci-wheel:-cuda12-azp-1` | `rapidsai/ci-wheel:-cuda-rockylinux8-py3.11`| libucxx/ucxx wheels (CUDA 12) | + +`` is the RAPIDS CalVer, e.g. `26.08`. The wrappers only add a `chmod` so the +non-root UID Azure runs steps as can write (`/opt/conda`, `/pyenv`) and `gdb` for +stack capture. Each image is a multi-arch manifest (`linux/amd64` + `linux/arm64`); +Azure picks the right arch per agent automatically. + +## When to rebuild + +RAPIDS ships ~every 2 months (CalVer `YY.MM`). We do **not** bump on every release: +the ucxx repo is pinned to a specific `main` commit, so we rebuild these images only +when we intentionally advance that pin to a commit whose base images/scripts require +a newer RAPIDS version. The CI images, the ucxx ref, and the in-repo CI scripts move +together. + +## Prerequisites + +- Two builder hosts, one `x86_64` and one `aarch64`, each with Docker + `buildx`. +- The host you run on must key-auth (non-interactive) ssh to the other host - the + remote buildkit runs over that ssh context. One direction suffices: run the + script from whichever host holds the key to the other. +- Harbor `ucx/` push credentials (`HARBOR_UCX_USER` / `HARBOR_UCX_PASSWORD`). + +(The specific builder hosts and credential source are environment-specific and are +not committed here - supply them via the environment, see below.) + +## Procedure + +### 1. Confirm the new RAPIDS base tags exist + +The build script preflights this; to check by hand (no login needed): + +```sh +VER=26.08 +docker manifest inspect rapidsai/ci-conda:${VER}-latest +docker manifest inspect rapidsai/ci-wheel:${VER}-cuda-rockylinux8-py3.11 +docker manifest inspect rapidsai/ci-wheel:${VER}-cuda-rockylinux8-py3.11 +``` + +The exact CUDA (``/``) and Python versions live in ONE place: +the `BASE_IMAGE` args of `docker-compose-rapidsai.yml` (the build script derives +its preflight from them). If RAPIDS moved a CUDA or Python version, update the +compose **and** the matching `rapids_cuda_version` / `rapids_py_version` in +`pr/main.yml` + `azure-pipelines-release.yml`. + +### 2. Build + push all three images, multi-arch + +One command, run from either builder (the script detects the host arch and appends +the opposite arch over ssh). Set the opposite-arch node endpoint and harbor creds in +the environment: + +```sh +RAPIDS_VER=26.08 \ +ARM_NODE=ssh://@ \ # when running on an x86 host +HARBOR_UCX_USER=... HARBOR_UCX_PASSWORD=... \ +./rapidsai-images-build.sh +``` + +The script preflights the base tags, creates a 2-node buildx builder (each node +pinned to its native arch so buildx never falls back to slow QEMU emulation), +logs in to harbor, runs `docker buildx bake --push` (builds both arches on their +native node and writes the manifest in one shot), logs out, and verifies each +pushed tag is multi-arch. + +Credential hygiene: pass creds via the environment, never leave them at rest on a +builder. The script logs out on exit. + +### 3. Point the pipelines at the new images + +Bump the old `` to the new one in the only places that reference it: + +- `buildlib/dockers/rapidsai-ci-conda.Dockerfile` (`BASE_IMAGE`) +- `buildlib/dockers/rapidsai-ci-wheel.Dockerfile` (`BASE_IMAGE`) +- `buildlib/pr/main.yml` (image tags) +- `buildlib/azure-pipelines-release.yml` (image tags) + +```sh +grep -rn buildlib/dockers buildlib/pr/main.yml buildlib/azure-pipelines-release.yml +python3 -c "import yaml; yaml.safe_load(open('buildlib/pr/main.yml'))" +python3 -c "import yaml; yaml.safe_load(open('buildlib/azure-pipelines-release.yml'))" +``` + +Commit on the PR-pipeline branch, rebase the release-pipeline branch onto it (so it +inherits the Dockerfile + `main.yml` bump), bump the release.yml tags there, push, +and re-run CI. + +## Why multi-arch + buildx + +The legacy UCX release images (`docker-compose-.yml` + `push-release-images.sh`) +build each arch separately and push to per-arch paths (`ucx/x86_64/...`, +`ucx/aarch64/...`), predating good multi-arch tooling. The RAPIDS wrappers use the +modern path: one compose file with `build.platforms`, built via `docker buildx bake` +against a 2-node native builder, producing a single multi-arch manifest per image - +fewer pipeline container resources (one tag, arch auto-resolved) and one command +instead of a per-arch loop plus manual `imagetools create`. + +## Troubleshooting + +- **only one architecture in the pushed manifest** - `docker buildx bake` does not + honor the compose `build.platforms` key; the script forces both platforms with + `--set '*.platform=...'`. If you bake by hand, pass it too. +- **`runc run failed ... can't mask dir /proc/acpi`** - the floating + `moby/buildkit:buildx-stable-1` image ships a runc that fails on the builders' + older kernels; the script pins the buildkit image (`BUILDKIT_IMAGE`). +- **`no space left on device` under `/var/lib/buildkit`** - the buildx builder's + own container volume filled up. `docker buildx rm ` and rerun (the + script recreates it). +- **arm64 build extremely slow** - it landed on the x86 node's QEMU emulation instead + of the native arm node. The `--platform` pins in the script prevent this; if you + build by hand, pin each node to one arch. +- **`docker context create` / ssh fails** - the host you run on can't key-auth to the + other node. Fix ssh keys first; the remote buildkit runs over that ssh context. +- **manifest shows only one arch after a push** - a node's push didn't complete (seen + when a backgrounded ssh build was killed early). Re-push the missing arch image and + re-create the manifest (`docker buildx imagetools create`), or just re-run the bake. +- **`docker login` 401 on `ucx/` repos** - the cached login is for a different harbor + project. Re-login with `HARBOR_UCX_USER`.