BUILD: add RAPIDS wrapper-image build tooling

buildlib/dockers/docker-compose-rapidsai.yml

@@ -0,0 +1,44 @@
+# Multi-arch (amd64 + arm64) build of the UCXX rapidsai CI wrapper images.
+#
+# Driven by build-rapidsai-images.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.08 docker buildx bake --builder <multi> -f docker-compose-rapidsai.yml --push
+#
+# Bumping RAPIDS: set RAPIDS_VER. If RAPIDS moves the wheel CUDA/py tag
+# (e.g. cuda13.2.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: harbor.mellanox.com/ucx/rapidsai-ci-conda:${RAPIDS_VER}-azp-1
+    build:
+      context: .
+      dockerfile: rapidsai-ci-conda.Dockerfile
+      platforms:
+        - linux/amd64
+        - linux/arm64
+      args:
+        BASE_IMAGE: rapidsai/ci-conda:${RAPIDS_VER}-latest
+
+  rapidsai-ci-wheel-cuda13:
+    image: harbor.mellanox.com/ucx/rapidsai-ci-wheel:${RAPIDS_VER}-cuda13-azp-1
+    build:
+      context: .
+      dockerfile: rapidsai-ci-wheel.Dockerfile
+      platforms:
+        - linux/amd64
+        - linux/arm64
+      args:
+        BASE_IMAGE: rapidsai/ci-wheel:${RAPIDS_VER}-cuda13.2.0-rockylinux8-py3.11
+
+  rapidsai-ci-wheel-cuda12:
+    image: harbor.mellanox.com/ucx/rapidsai-ci-wheel:${RAPIDS_VER}-cuda12-azp-1
+    build:
+      context: .
+      dockerfile: rapidsai-ci-wheel.Dockerfile
+      platforms:
+        - linux/amd64
+        - linux/arm64
+      args:
+        BASE_IMAGE: rapidsai/ci-wheel:${RAPIDS_VER}-cuda12.9.1-rockylinux8-py3.11

buildlib/dockers/rapidsai-images-build.sh

@@ -0,0 +1,84 @@
+#!/bin/bash -eE
+#
+# 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=<ver> \
+#   ARM_NODE=ssh://<user>@<arm-builder>   # required when running on an x86 host
+#   X86_NODE=ssh://<user>@<x86-builder>   # 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)
+#   BUILDER      buildx builder name (default: ucx-rapidsai)
+#   HARBOR_UCX_USER / HARBOR_UCX_PASSWORD   harbor push credentials (required)
+
+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}
+
+# 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://<user>@<arm-builder>} ;;
+  aarch64) LOCAL_PLAT=linux/arm64; REMOTE_PLAT=linux/amd64
+           REMOTE_NODE=${X86_NODE:?set X86_NODE=ssh://<user>@<x86-builder>} ;;
+  *) echo "ERROR: unexpected host arch $(uname -m)" >&2; exit 1 ;;
+esac
+
+# 1. preflight - the rapidsai base tags for this version must exist
+for t in "rapidsai/ci-conda:${RAPIDS_VER}-latest" \
+         "rapidsai/ci-wheel:${RAPIDS_VER}-cuda13.2.0-rockylinux8-py3.11" \
+         "rapidsai/ci-wheel:${RAPIDS_VER}-cuda12.9.1-rockylinux8-py3.11"; do
+  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
+
+# 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).
+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 \
+    --platform "$LOCAL_PLAT" --bootstrap                                                # local (host arch)
+  docker buildx create --name "$BUILDER" --append --node "${BUILDER}-remote" \
+    --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
+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" --push
+
+# 4. verify the 3 manifests are multi-arch
+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"
+  docker manifest inspect "$REGISTRY/ucx/$img" 2>/dev/null \
+    | grep -o '"architecture": "[a-z0-9]*"' | tr '\n' ' '
+  echo
+done
+echo "DONE: rapidsai-ci-{conda,wheel cuda12/13}:${RAPIDS_VER}-azp-1 (amd64+arm64)"

buildlib/dockers/rapidsai-images.md

@@ -0,0 +1,116 @@
+# 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 Dockerfiles,
+a compose file, and a build script to (re)build and push them.
+
+| Wrapper image (`ucx/`)                     | RAPIDS base                                            | Used for                   |
+|--------------------------------------------|--------------------------------------------------------|----------------------------|
+| `rapidsai-ci-conda:<VER>-azp-1`            | `rapidsai/ci-conda:<VER>-latest`                       | conda build + tests        |
+| `rapidsai-ci-wheel:<VER>-cuda13-azp-1`     | `rapidsai/ci-wheel:<VER>-cuda13.2.0-rockylinux8-py3.11`| libucxx/ucxx wheels (CUDA 13) |
+| `rapidsai-ci-wheel:<VER>-cuda12-azp-1`     | `rapidsai/ci-wheel:<VER>-cuda12.9.1-rockylinux8-py3.11`| libucxx/ucxx wheels (CUDA 12) |
+
+`<VER>` 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. Both directions work.
+- 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}-cuda13.2.0-rockylinux8-py3.11
+docker manifest inspect rapidsai/ci-wheel:${VER}-cuda12.9.1-rockylinux8-py3.11
+```
+
+If RAPIDS moved the CUDA or Python in the wheel tag (e.g. `cuda13.2.0` -> `13.3.0`,
+or `py3.11` -> `py3.12`), update the two wheel `BASE_IMAGE` args in
+`docker-compose-rapidsai.yml` **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://<user>@<arm-builder> \   # 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 `<VER>` 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 <VER> 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-<arch>.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
+
+- **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`.