Use RAPIDS_BRANCH for RAPIDS CI clones

.github/workflows/build-all-rapids-repos.yml

@@ -103,7 +103,7 @@ jobs:
 
         # Clone all the repos
         run_command "Clone RAPIDS repositories" \
-          clone-all -j$(nproc) -v -q --clone-upstream --depth 1 --single-branch --shallow-submodules --no-update-env;
+          clone-all -j$(nproc) -v -q --branch "$(< RAPIDS_BRANCH)" --clone-upstream --depth 1 --single-branch --shallow-submodules --no-update-env;
 
         run_command "Create RAPIDS python environment" \
           rapids-post-start-command;

RAPIDS_BRANCH

@@ -0,0 +1 @@
+release/26.06

ci/release/update-version.sh

@@ -81,13 +81,16 @@ NEXT_FULL_TAG_PEP440=$(python -c "from packaging.version import Version; print(V
 
 # Log update context
 if [[ "${RUN_CONTEXT}" == "main" ]]; then
+    RAPIDS_BRANCH_NAME="main"
     echo "Preparing development branch update $CURRENT_TAG => $NEXT_FULL_TAG (targeting main branch)"
 elif [[ "${RUN_CONTEXT}" == "release" ]]; then
+    RAPIDS_BRANCH_NAME="release/${NEXT_SHORT_TAG}"
     echo "Preparing release branch update $CURRENT_TAG => $NEXT_FULL_TAG (targeting release/${NEXT_SHORT_TAG} branch)"
 fi
 
 # Centralized version file update
 echo "${NEXT_FULL_TAG}" > VERSION
+echo "${RAPIDS_BRANCH_NAME}" > RAPIDS_BRANCH
 
 # Inplace sed replace; workaround for Linux and Mac
 function sed_runner() {

features/src/rapids-build-utils/install.sh

@@ -56,7 +56,7 @@ cp -ar ./opt/rapids-build-utils /opt/;
 
 declare -a commands=(
     checkout-same-branch
-    maybe-clean-build-dir
+    convert-ucxx-branch
     generate-docstring
     generate-scripts
     get-cmake-build-dir
@@ -68,6 +68,7 @@ declare -a commands=(
     make-pip-env
     make-vscode-workspace
     merge-compile-commands-json
+    maybe-clean-build-dir
     post-attach-command
     post-start-command
     pull-repositories

features/src/rapids-build-utils/opt/rapids-build-utils/bin/checkout-same-branch.sh

@@ -17,25 +17,6 @@
 # shellcheck disable=SC1091
 . rapids-generate-docstring;
 
-# Function to convert ucx branch names
-convert-ucx-branch() {
-    repo=$1
-    custom_branch=$2
-    normalized_branch="${custom_branch}"
-
-    if [[ "${repo}" == "ucx"* ]]; then
-        # Only convert branches that match the pattern release/YY.MM
-        if [[ "${custom_branch}" =~ ^release/[0-9]{2}\.[0-9]{2}$ ]]; then
-            RAPIDS_VERSION=$(echo "${custom_branch}" | awk '{split($0, a, "/"); print a[2]}')
-            # Get UCX version associated w/ RAPIDS version
-            UCX_VERSION="$(curl -sL https://version.gpuci.io/rapids/${RAPIDS_VERSION})"
-            normalized_branch="release/${UCX_VERSION}"
-        fi
-    fi
-
-    echo "${normalized_branch}"
-}
-
 checkout_same_branch() {
     local -;
     set -euo pipefail;
@@ -68,7 +49,7 @@ checkout_same_branch() {
         local path="${!repo_path}";
 
         # Skip ucxx repositories when determining common branches
-        if [[ "${name}" == "ucx"* ]]; then
+        if [[ "${name}" == "ucxx" ]]; then
             continue;
         fi
 
@@ -155,10 +136,10 @@ checkout_same_branch() {
 
         # Apply normalization for ucxx repositories
         local repo_name_val="${!repo_name}";
-        local normalized_branch;
-        normalized_branch=$(convert-ucx-branch "${repo_name_val}" "${branch}");
 
-        if [[ "${repo_name_val}" == "ucx"* ]]; then
+        if [[ "${repo_name_val}" == "ucxx" ]]; then
+            local normalized_branch;
+            normalized_branch=$(rapids-convert-ucxx-branch "${branch}");
             echo "Using normalized branch '${normalized_branch}' for repository '${repo_name_val}'";
             branch="${normalized_branch}";
         fi

features/src/rapids-build-utils/opt/rapids-build-utils/bin/convert-ucxx-branch.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+# Usage:
+#  rapids-convert-ucxx-branch <branch>
+#
+# Convert a RAPIDS branch name to the corresponding UCXX branch name.
+
+convert_ucxx_branch() {
+    local branch="${1:?fatal: missing required positional argument <branch>}"
+
+    if [[ "${branch}" =~ ^release/[0-9]{2}\.[0-9]{2}$ ]]; then
+        local rapids_version="${branch#release/}"
+        branch="release/$(curl -sL "https://version.gpuci.io/rapids/${rapids_version}")"
+    fi
+
+    echo "${branch}"
+}
+
+convert_ucxx_branch "$@" <&0

features/src/rapids-build-utils/opt/rapids-build-utils/bin/tmpl/all.clone.tmpl.sh

@@ -11,6 +11,7 @@
 #                               (default: false)
 #
 # Options that require values:
+#  -b,--branch <branch_or_tag>  Check the repos out to <branch_or_tag>.
 #  -j,--parallel <num>          Clone <num> repos in parallel.
 #                               (default: 1)
 
@@ -21,7 +22,7 @@ clone_all() {
     local -;
     set -euo pipefail;
 
-    eval "$(_parse_args --take '-j,--parallel -v,--verbose --no-update-env' "$@" <&0)";
+    eval "$(_parse_args --take '-b,--branch -j,--parallel -v,--verbose --no-update-env' "$@" <&0)";
 
     eval "$(rapids-get-num-archs-jobs-and-load --archs 1 --max-device-obj-memory-usage 1 "$@")";
 
@@ -36,7 +37,22 @@ clone_all() {
     echo ${NAMES} \
   | tr '[:space:]' '\0' \
   | xargs ${v:+-t} ${_o} -r -0 -P${n_jobs} -I% bash -c \
-  " if command -V clone-% >/dev/null 2>&1; then if ! clone-% -j ${n_arch} --no-update-env ${OPTS[*]@Q} ${v[*]@Q}; then exit 255; fi; fi";
+  '
+    repo="$1";
+    branch="$2";
+    n_arch="$3";
+    shift 3;
+    branch_args=();
+    if test -n "${branch:+x}"; then
+        if [[ "${repo}" == "ucxx" ]]; then
+            branch="$(rapids-convert-ucxx-branch "${branch}")";
+        fi;
+        branch_args=(--branch "${branch}");
+    fi;
+    if command -V "clone-${repo}" >/dev/null 2>&1; then
+        if ! "clone-${repo}" -j "${n_arch}" --no-update-env "${branch_args[@]}" "$@"; then exit 255; fi;
+    fi
+  ' _ % "${branch[0]:-}" "${n_arch}" "${OPTS[@]}" "${v[@]}";
 
     if ! test -n "${no_update_env:+x}"; then
         rapids-post-start-command;