merge in cibuildwheels-imgs repo

.github/workflows/build-and-publish-images.yml

@@ -46,6 +46,7 @@ jobs:
       PYTHON_VER: ${{ matrix.PYTHON_VER }}
       IMAGE_REPO: ${{ matrix.IMAGE_REPO }}
       IMAGE_NAME: ${{ matrix.IMAGE_NAME }}
+      MANYLINUX_VER: ${{ matrix.MANYLINUX_VER }}
 
   build-multiarch-manifest:
     name: manifest (${{ matrix.CUDA_VER }}, ${{ matrix.PYTHON_VER }}, ${{ matrix.LINUX_VER }}, ${{ matrix.IMAGE_REPO }})

.github/workflows/build-image.yml

@@ -21,6 +21,9 @@ on:
       IMAGE_NAME:
         required: true
         type: string
+      MANYLINUX_VER:
+        required: false
+        type: string
 
 jobs:
   run:
@@ -49,15 +52,22 @@ jobs:
         with:
           driver: docker
           endpoint: builders
+      - name: Generate Build Args
+        id: generate-build-args
+        run: ci/compute-build-args.sh
+        env:
+          IMAGE_REPO: ${{ inputs.IMAGE_REPO }}
+          CUDA_VER: ${{ inputs.CUDA_VER }}
+          LINUX_VER: ${{ inputs.LINUX_VER }}
+          PYTHON_VER: ${{ inputs.PYTHON_VER }}
+          ARCH: ${{ matrix.ARCH }}
       - name: Build image
         uses: docker/build-push-action@v4
         with:
           context: context
-          file: Dockerfile
+          file: dockerfiles/${{ inputs.IMAGE_REPO }}.Dockerfile
           push: true
           pull: true
           build-args: |
-            CUDA_VER=${{ inputs.CUDA_VER }}
-            LINUX_VER=${{ inputs.LINUX_VER }}
-            PYTHON_VER=${{ inputs.PYTHON_VER }}
+            ${{ steps.generate-build-args.outputs.ARGS }}
           tags: ${{ inputs.IMAGE_NAME }}-${{ matrix.ARCH }}

ci/compute-build-args.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+set -euo pipefail
+
+ARGS=""
+case "${IMAGE_REPO}" in
+  ci)
+    ARGS="CUDA_VER=${CUDA_VER}
+LINUX_VER=${LINUX_VER}
+PYTHON_VER=${PYTHON_VER}"
+    ;;
+  *)
+    MANYLINUX_VER="manylinux_2_17"
+    if [[
+        "${LINUX_VER}" == "ubuntu18.04" ||
+        "${LINUX_VER}" == "ubuntu20.04"
+    ]]; then
+        MANYLINUX_VER="manylinux_2_31"
+    fi
+    ARGS="CUDA_VER=${CUDA_VER}
+LINUX_VER=${LINUX_VER}
+PYTHON_VER=${PYTHON_VER}
+CPU_ARCH=${ARCH}
+REAL_ARCH=$(arch)
+MANYLINUX_VER=${MANYLINUX_VER}"
+    ;;
+esac
+
+echo "ARGS<<EOF" >> "$GITHUB_OUTPUT"
+echo "$ARGS" >> "$GITHUB_OUTPUT"
+echo "EOF" >> "$GITHUB_OUTPUT"

ci/compute-latest-versions.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Computes versions used for "latest" tag based on "matrix.yaml"
+# Computes versions used for "latest" tag based on "matrices/ci-matrix.yaml"
 # values. Will also check to ensure that the "latest" values are
 # included in the matrix values.
 # Example Usage:
@@ -11,14 +11,14 @@ export CUDA_KEY="CUDA_VER"
 export PYTHON_KEY="PYTHON_VER"
 
 # Get latest values
-LATEST_LINUX_VER=$(yq '.LATEST_VERSIONS.[strenv(LINUX_KEY)]' matrix.yaml)
-LATEST_CUDA_VER=$(yq '.LATEST_VERSIONS.[strenv(CUDA_KEY)]' matrix.yaml)
-LATEST_PYTHON_VER=$(yq '.LATEST_VERSIONS.[strenv(PYTHON_KEY)]' matrix.yaml)
+LATEST_LINUX_VER=$(yq '.LATEST_VERSIONS.[strenv(LINUX_KEY)]' matrices/ci-matrix.yaml)
+LATEST_CUDA_VER=$(yq '.LATEST_VERSIONS.[strenv(CUDA_KEY)]' matrices/ci-matrix.yaml)
+LATEST_PYTHON_VER=$(yq '.LATEST_VERSIONS.[strenv(PYTHON_KEY)]' matrices/ci-matrix.yaml)
 
 # Get matrix array values
-LINUX_VERS=$(yq '.[strenv(LINUX_KEY)]' matrix.yaml)
-CUDA_VERS=$(yq '.[strenv(CUDA_KEY)]' matrix.yaml)
-PYTHON_VERS=$(yq '.[strenv(PYTHON_KEY)]' matrix.yaml)
+LINUX_VERS=$(yq '.[strenv(LINUX_KEY)]' matrices/ci-matrix.yaml)
+CUDA_VERS=$(yq '.[strenv(CUDA_KEY)]' matrices/ci-matrix.yaml)
+PYTHON_VERS=$(yq '.[strenv(PYTHON_KEY)]' matrices/ci-matrix.yaml)
 
 # Ensure matrix array values contain latest values
 for KEY in "${LINUX_KEY}" "${CUDA_KEY}" "${PYTHON_KEY}"; do

ci/compute-matrix.sh

@@ -13,4 +13,12 @@ case "${BUILD_TYPE}" in
     ;;
 esac
 
-yq -o json '. | del(.LATEST_VERSIONS)' matrix.yaml | jq -c 'include "ci/compute-matrix"; compute_matrix(.)'
+CI_MATRIX=$(yq -o json '. | del(.LATEST_VERSIONS)' matrices/ci-matrix.yaml | jq -c 'include "ci/compute-ci-matrix"; compute_matrix(.)')
+WHEELS_MATRIX=$(yq -o json matrices/wheels-matrix.yaml | jq -c 'include "ci/compute-wheels-matrix"; compute_matrix(.)')
+
+COMBINED_MATRIX=$(jq -c -n \
+  --argjson ci_matrix "$CI_MATRIX" \
+  --argjson wheels_matrix "$WHEELS_MATRIX" \
+  '{"include": ($ci_matrix.include + $wheels_matrix.include)}')
+
+echo "$COMBINED_MATRIX"

ci/compute-wheels-matrix.jq

@@ -0,0 +1,73 @@
+def compute_arch($x):
+  ["amd64"] |
+  if
+    ["ubuntu18.04", "centos7"] | index($x.LINUX_VER) != null
+  then
+    .
+  else
+    . + ["arm64"]
+  end |
+  $x + {ARCHES: .};
+
+# Checks the current entry to see if it matches the given exclude
+def matches($entry; $exclude):
+  all($exclude | to_entries | .[]; $entry[.key] == .value);
+
+def compute_repo($x):
+  if
+    env.BUILD_TYPE == "pull-request"
+  then
+    "staging"
+  else
+    $x.IMAGE_REPO
+  end;
+
+def compute_tag_prefix($x):
+  if
+    env.BUILD_TYPE == "branch"
+  then
+    ""
+  else
+    $x.IMAGE_REPO + "-" + env.PR_NUM + "-"
+  end;
+
+def compute_manylinux_version($x):
+  if
+    ["ubuntu18.04", "ubuntu20.04"] | index($x.LINUX_VER) != null
+  then
+    "manylinux_2_31"
+  else
+    "manylinux_2_17"
+  end |
+  $x + {MANYLINUX_VER: .};
+
+def compute_image_name($x):
+  compute_repo($x) as $repo |
+  compute_tag_prefix($x) as $tag_prefix |
+  "rapidsai/" + $repo + ":" + $tag_prefix + "cuda" + $x.CUDA_VER + "-" + $x.LINUX_VER + "-" + "py" + $x.PYTHON_VER |
+  $x + {IMAGE_NAME: .};
+
+
+# Checks the current entry to see if it matches any of the excludes.
+# If so, produce no output. Otherwise, output the entry.
+def filter_excludes($entry; $excludes):
+  select(any($excludes[]; matches($entry; .)) | not);
+
+def lists2dict($keys; $values):
+  reduce range($keys | length) as $ind ({}; . + {($keys[$ind]): $values[$ind]});
+
+def compute_matrix($input):
+  ($input.exclude // []) as $excludes |
+  $input | del(.exclude) |
+  keys_unsorted as $matrix_keys |
+  to_entries |
+  map(.value) |
+  [
+    combinations |
+    lists2dict($matrix_keys; .) |
+    filter_excludes(.; $excludes) |
+    compute_arch(.) |
+    compute_manylinux_version(.) |
+    compute_image_name(.)
+  ] |
+  {include: .};

ci/create-multiarch-manifest.sh

@@ -1,6 +1,13 @@
 #!/bin/bash
 set -euo pipefail
 
+# if IMAGE_REPO is "ci" then compute the latest versions below
+if [[ "${IMAGE_REPO}" != "ci" ]]; then
+  LATEST_CUDA_VER=$(yq '.CUDA_VER | sort | .[-1]' matrices/wheels-matrix.yaml)
+  LATEST_PYTHON_VER=$(yq -o json '.PYTHON_VER' matrices/wheels-matrix.yaml | jq -r 'max_by(split(".") | map(tonumber))')
+  LATEST_UBUNTU_VER=$(yq '.LINUX_VER | map(select(. == "*ubuntu*")) | sort | .[-1]' matrices/wheels-matrix.yaml)
+fi
+
 source_tags=()
 tag="${IMAGE_NAME}"
 for arch in $(echo "${ARCHES}" | jq .[] -r); do

context/citestwheel.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+
+set -eoxu pipefail
+
+export RAPIDS_PY_WHEEL_NAME="${RAPIDS_PY_WHEEL_NAME:-}"
+export RAPIDS_PY_VERSION="${RAPIDS_PY_VERSION:-}"
+export CIBW_TEST_EXTRAS="${CIBW_TEST_EXTRAS:-}"
+export CIBW_TEST_COMMAND="${CIBW_TEST_COMMAND:-}"
+export RAPIDS_BEFORE_TEST_COMMANDS_AMD64="${RAPIDS_BEFORE_TEST_COMMANDS_AMD64:-}"
+export RAPIDS_BEFORE_TEST_COMMANDS_ARM64="${RAPIDS_BEFORE_TEST_COMMANDS_ARM64:-}"
+export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-}"
+
+mkdir -p ./dist
+
+arch=$(uname -m)
+
+# need this to init pyenv first
+eval "$(pyenv init -)"
+
+# use pyenv to set appropriate python as default before citestwheel
+pyenv global "${RAPIDS_PY_VERSION}" && python --version
+
+rapids-download-wheels-from-s3 ./dist
+
+if [ "${arch}" == "x86_64" ]; then
+        sh -c "${RAPIDS_BEFORE_TEST_COMMANDS_AMD64}"
+elif [ "${arch}" == "aarch64" ]; then
+        sh -c "${RAPIDS_BEFORE_TEST_COMMANDS_ARM64}"
+fi
+
+# see: https://cibuildwheel.readthedocs.io/en/stable/options/#test-extras
+extra_requires_suffix=''
+if [ "${CIBW_TEST_EXTRAS}" != "" ]; then
+        extra_requires_suffix="[${CIBW_TEST_EXTRAS}]"
+fi
+
+# echo to expand wildcard before adding `[extra]` requires for pip
+python -m pip install --verbose $(echo ./dist/${RAPIDS_PY_WHEEL_NAME}*.whl)$extra_requires_suffix
+
+python -m pip check
+
+sh -c "${CIBW_TEST_COMMAND}"

dockerfiles/ci-wheel.Dockerfile

@@ -0,0 +1,145 @@
+ARG CUDA_VER=11.8.0
+ARG LINUX_VER=ubuntu20.04
+ARG REAL_ARCH=x86_64
+
+ARG BASE_IMAGE=nvidia/cuda:${CUDA_VER}-devel-${LINUX_VER}
+FROM ${BASE_IMAGE}
+
+ARG CUDA_VER
+ARG LINUX_VER
+ARG CPU_ARCH
+ARG REAL_ARCH
+ARG PYTHON_VER=3.9
+ARG MANYLINUX_VER
+ARG POLICY=${MANYLINUX_VER}
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Set RAPIDS versions env variables
+ENV RAPIDS_CUDA_VERSION="${CUDA_VER}"
+ENV RAPIDS_PY_VERSION="${PYTHON_VER}"
+
+# RAPIDS pip index
+ENV PIP_EXTRA_INDEX_URL="https://pypi.k8s.rapids.ai/simple"
+
+ENV PYENV_ROOT="/pyenv"
+ENV PATH="/pyenv/bin:/pyenv/shims:$PATH"
+
+RUN case "${LINUX_VER}" in \
+    "ubuntu"*) \
+        apt update -y && apt install -y jq build-essential software-properties-common wget gcc zlib1g-dev libbz2-dev libssl-dev libreadline-dev libsqlite3-dev libffi-dev curl git libncurses5-dev libnuma-dev openssh-client libcudnn8-dev zip libopenblas-dev liblapack-dev protobuf-compiler autoconf automake libtool cmake && rm -rf /var/lib/apt/lists/* \
+        && add-apt-repository ppa:git-core/ppa && add-apt-repository ppa:ubuntu-toolchain-r/test && apt update -y && apt install -y git gcc-9 g++-9 && add-apt-repository -r ppa:git-core/ppa && add-apt-repository -r ppa:ubuntu-toolchain-r/test \
+        && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 90 --slave /usr/bin/g++ g++ /usr/bin/g++-9 --slave /usr/bin/gcov gcov /usr/bin/gcov-9 \
+      ;; \
+    "centos"*) \
+        yum update --exclude=libnccl* -y && yum install -y epel-release wget gcc zlib-devel bzip2 bzip2-devel readline-devel sqlite sqlite-devel xz xz-devel libffi-devel curl git ncurses-devel numactl numactl-devel openssh-clients libcudnn8-devel zip blas-devel lapack-devel protobuf-compiler autoconf automake libtool centos-release-scl scl-utils cmake && yum clean all \
+        && yum remove -y git && yum install -y https://packages.endpointdev.com/rhel/7/os/x86_64/endpoint-repo.x86_64.rpm && yum install -y git jq devtoolset-11 && yum remove -y endpoint-repo \
+        && echo -e ' \
+        #!/bin/bash\n \
+        source scl_source enable devtoolset-11\n \
+        ' > /etc/profile.d/enable_devtools.sh \
+        && pushd tmp \
+        && wget https://ftp.openssl.org/source/openssl-1.1.1k.tar.gz \
+        && tar -xzvf openssl-1.1.1k.tar.gz \
+        && cd openssl-1.1.1k \
+        && ./config --prefix=/usr --openssldir=/etc/ssl --libdir=lib no-shared zlib-dynamic \
+        && make \
+        && make install \
+        && popd \
+      ;; \
+    *) \
+      echo "Unsupported LINUX_VER: ${LINUX_VER}" && exit 1; \
+      ;; \
+  esac
+
+# Download and install GH CLI tool v2.32.0
+ARG GH_VERSION=2.32.0
+RUN <<EOF
+set -e
+wget https://github.com/cli/cli/releases/download/v${GH_VERSION}/gh_${GH_VERSION}_linux_${CPU_ARCH}.tar.gz
+tar -xf gh_*.tar.gz
+mv gh_*/bin/gh /usr/local/bin
+rm -rf gh_*
+EOF
+
+# Install sccache
+ARG SCCACHE_VERSION=0.5.0
+
+RUN curl -o /tmp/sccache.tar.gz \
+        -L "https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-"${REAL_ARCH}"-unknown-linux-musl.tar.gz" && \
+        tar -C /tmp -xvf /tmp/sccache.tar.gz && \
+        mv "/tmp/sccache-v${SCCACHE_VERSION}-"${REAL_ARCH}"-unknown-linux-musl/sccache" /usr/bin/sccache && \
+        chmod +x /usr/bin/sccache
+
+# Set AUDITWHEEL_* env vars for use with auditwheel
+ENV AUDITWHEEL_POLICY=${POLICY} AUDITWHEEL_ARCH=${REAL_ARCH} AUDITWHEEL_PLAT=${POLICY}_${REAL_ARCH}
+
+# Set sccache env vars
+ENV CMAKE_CUDA_COMPILER_LAUNCHER=sccache
+ENV CMAKE_CXX_COMPILER_LAUNCHER=sccache
+ENV CMAKE_C_COMPILER_LAUNCHER=sccache
+ENV SCCACHE_BUCKET=rapids-sccache-east
+ENV SCCACHE_REGION=us-east-2
+ENV SCCACHE_IDLE_TIMEOUT=32768
+ENV SCCACHE_S3_USE_SSL=true
+ENV SCCACHE_S3_NO_CREDENTIALS=false
+
+# Install ucx
+ARG UCX_VERSION=1.14.1
+RUN mkdir -p /ucx-src && cd /ucx-src &&\
+    git clone https://github.com/openucx/ucx -b v${UCX_VERSION} ucx-git-repo &&\
+    cd ucx-git-repo && \
+    ./autogen.sh && \
+    ./contrib/configure-release \
+       --prefix=/usr               \
+       --enable-mt                 \
+       --enable-cma                \
+       --enable-numa               \
+       --with-gnu-ld               \
+       --with-sysroot              \
+       --without-verbs             \
+       --without-rdmacm            \
+       --with-cuda=/usr/local/cuda && \
+    CPPFLAGS=-I/usr/local/cuda/include make -j && \
+    make install && \
+    cd / && \
+    rm -rf /ucx-src/
+
+# Install pyenv
+RUN curl https://pyenv.run | bash
+
+# Create pyenvs
+# TODO: Determine if any cleanup of the pyenv layers is needed to shrink the container
+RUN pyenv update
+
+RUN case "${LINUX_VER}" in \
+    "ubuntu"*) \
+        pyenv install --verbose "${RAPIDS_PY_VERSION}" \
+      ;; \
+    "centos"*) \
+        # Need to specify the openssl location because of the install from source
+        CPPFLAGS="-I/usr/include/openssl" LDFLAGS="-L/usr/lib" pyenv install --verbose "${RAPIDS_PY_VERSION}" \
+      ;; \
+    *) \
+      echo "Unsupported LINUX_VER: ${LINUX_VER}" && exit 1; \
+      ;; \
+  esac
+
+RUN pyenv global ${PYTHON_VER} && python -m pip install auditwheel patchelf twine && pyenv rehash
+
+# Install latest gha-tools
+RUN wget https://github.com/rapidsai/gha-tools/releases/latest/download/tools.tar.gz -O - | tar -xz -C /usr/local/bin
+
+# Install the AWS CLI
+RUN mkdir -p /aws_install && cd /aws_install && \
+    curl "https://s3.amazonaws.com/aws-cli/awscli-bundle.zip" -o "awscli-bundle.zip" && \
+    unzip awscli-bundle.zip && \
+    ./awscli-bundle/install -i /usr/local/aws -b /usr/local/bin/aws && \
+    cd / && \
+    rm -rf /aws_install
+
+# Mark all directories as safe for git so that GHA clones into the root don't
+# run into issues
+RUN git config --system --add safe.directory '*'
+
+CMD ["/bin/bash"]