BUILD: add RAPIDS wrapper-image build tooling

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 <multi> -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

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=<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); 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://<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 - 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)"

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:<VER>-azp-1`            | `rapidsai/ci-conda:<VER>-latest`                       | conda build + tests        |
+| `rapidsai-ci-wheel:<VER>-cuda13-azp-1`     | `rapidsai/ci-wheel:<VER>-cuda<CU13>-rockylinux8-py3.11`| libucxx/ucxx wheels (CUDA 13) |
+| `rapidsai-ci-wheel:<VER>-cuda12-azp-1`     | `rapidsai/ci-wheel:<VER>-cuda<CU12>-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. 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<CU13>-rockylinux8-py3.11
+docker manifest inspect rapidsai/ci-wheel:${VER}-cuda<CU12>-rockylinux8-py3.11
+```
+
+The exact CUDA (`<CU12>`/`<CU13>`) 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://<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
+
+- **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 <builder>` 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`.