Benchmark: make it possible to build container with UCX built from source (#306)

* Support building UCX from source

* Make nixl optional, just point to repo root

* Address comments

* fix indentation

Author: ovidiusm

benchmark/nixlbench/contrib/Dockerfile

@@ -16,10 +16,15 @@
 ARG BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
 ARG BASE_IMAGE_TAG="25.03-cuda12.8-devel-ubuntu24.04"
 
-FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG}
-
+# UCX argument is either "upstream" (default installed in base image) or "custom" (build from source)
+ARG UCX="upstream"
 ARG DEFAULT_PYTHON_VERSION="3.12"
 
+# --- Stage 1: Common OS setup ---
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS os_setup_stage
+
+# Re-declare for use in this stage
+ARG DEFAULT_PYTHON_VERSION
 RUN apt-get update -y && \
     DEBIAN_FRONTEND=noninteractive apt-get -y install \
     ninja-build \
@@ -34,6 +39,44 @@ RUN apt-get update -y && \
     etcd-server \
     etcd-client
 
+# --- Stage 2a: Represents using UCX from the base image ---
+FROM os_setup_stage AS ucx_upstream_image
+RUN echo "INFO: Using UCX from base image (UCX=${UCX})."
+
+# --- Stage 2b: Represents building UCX from source ---
+FROM os_setup_stage AS ucx_custom_image
+RUN mkdir -p /workspace/ucx
+COPY --from=ucx . /workspace/ucx
+
+RUN echo "INFO: Starting custom UCX build..." && \
+    apt-get update -y && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        autoconf automake libtool pkg-config make g++ \
+        libnuma-dev librdmacm-dev ibverbs-providers && \
+    echo "INFO: Removing pre-existing UCX installations..." && \
+    rm -rf /usr/local/ucx /opt/hpcx/ucx && \
+    cd /workspace/ucx && \
+    ./autogen.sh && \
+    echo "INFO: Building UCX..." && \
+    ./contrib/configure-release --prefix=/usr/local/ucx \
+                                --with-cuda=/usr/local/cuda \
+                                --enable-mt \
+                                --without-go && \
+    make -j$(nproc) && \
+    make install && \
+    cd / && \
+    echo "INFO: Finished building and installing UCX to /usr/local/ucx."
+
+# --- Stage 3: UCX Image Selection ---
+# This stage selects the correct UCX image based on the UCX argument
+FROM ucx_${UCX}_image AS ucx_image
+
+# --- Stage 4: Final Image Assembly ---
+# Re-declare ARGs needed in this final stage
+ARG DEFAULT_PYTHON_VERSION
+ARG WHL_PYTHON_VERSIONS="3.12"
+ARG WHL_PLATFORM="manylinux_2_39_x86_64"
+
 WORKDIR /workspace
 RUN git clone https://github.com/etcd-cpp-apiv3/etcd-cpp-apiv3.git &&\
 	cd etcd-cpp-apiv3 && mkdir build && cd build && \
@@ -67,8 +110,6 @@ RUN echo "/usr/local/nixl/lib/x86_64-linux-gnu" > /etc/ld.so.conf.d/nixl.conf &&
     ldconfig
 
 # Create the wheel
-ARG WHL_PYTHON_VERSIONS="3.12"
-ARG WHL_PLATFORM="manylinux_2_39_x86_64"
 RUN IFS=',' read -ra PYTHON_VERSIONS <<< "$WHL_PYTHON_VERSIONS" && \
     for PYTHON_VERSION in "${PYTHON_VERSIONS[@]}"; do \
         uv build --wheel --out-dir /tmp/dist --python $PYTHON_VERSION; \

benchmark/nixlbench/contrib/build.sh

@@ -18,10 +18,12 @@
 SOURCE_DIR=$(dirname "$(readlink -f "$0")")
 BUILD_CONTEXT=$(dirname "$(readlink -f "$SOURCE_DIR")")
 BUILD_CONTEXT_ARGS=""
-NIXL_BUILD_CONTEXT_ARGS=""
+NIXL_SRC=$(readlink -f "${SOURCE_DIR}/../../..")
+NIXL_BUILD_CONTEXT_ARGS="--build-context nixl=$NIXL_SRC"
 NIXL_BENCH_BUILD_CONTEXT_ARGS="--build-context nixlbench=$BUILD_CONTEXT/"
-NIXL_SRC=""
 DOCKER_FILE="${SOURCE_DIR}/Dockerfile"
+UCX_SRC=""
+UCX_BUILD_CONTEXT_ARGS=""
 commit_id=$(git rev-parse --short HEAD)
 
 # Get latest TAG and add COMMIT_ID for dev
@@ -63,7 +65,7 @@ get_options() {
         --nixl)
             if [ "$2" ]; then
                 NIXL_BUILD_CONTEXT_ARGS="--build-context nixl=$2"
-		NIXL_SRC=$2
+                NIXL_SRC=$2
                 shift
             else
                 missing_requirement $1
@@ -97,6 +99,15 @@ get_options() {
                 missing_requirement $1
             fi
             ;;
+        --ucx)
+            if [ "$2" ]; then
+                UCX_SRC=$2
+                UCX_BUILD_CONTEXT_ARGS="--build-context ucx=$2 --build-arg UCX=custom"
+                shift
+            else
+                missing_requirement $1
+            fi
+            ;;
         --)
             shift
             break
@@ -118,7 +129,7 @@ get_options() {
         error "ERROR: --nixl <path to nixl source> is required"
     fi
 
-    BUILD_CONTEXT_ARGS="$NIXL_BUILD_CONTEXT_ARGS $NIXL_BENCH_BUILD_CONTEXT_ARGS"
+    BUILD_CONTEXT_ARGS="$NIXL_BUILD_CONTEXT_ARGS $NIXL_BENCH_BUILD_CONTEXT_ARGS $UCX_BUILD_CONTEXT_ARGS"
 
     VERSION=v$latest_tag.dev.$commit_id
     if [ -z "$TAG" ]; then
@@ -131,6 +142,7 @@ show_build_options() {
     echo ""
     echo "Building NIXLBench Image"
     echo "NIXL Source: ${NIXL_SRC}"
+    echo "UCX Source: ${UCX_SRC} (optional)"
     echo "Image Tag: ${TAG}"
     echo "Build Context: ${BUILD_CONTEXT}"
     echo "Build Context Args: ${BUILD_CONTEXT_ARGS}"
@@ -144,6 +156,7 @@ show_help() {
     echo "  [--base base image]"
     echo "  [--base-image-tag base image tag]"
     echo "  [--nixlbench path/to/nixlbench/source/dir]"
+    echo "  [--ucx path/to/ucx/source/dir]"
     echo "  [--no-cache disable docker build cache]"
     echo "  [--os [ubuntu24|ubuntu22] to select Ubuntu version]"
     echo "  [--python-versions python versions to build for, comma separated]"