Build and container fixes for aarch64 (#408)

* Fix gds_path default

CUDA provides symlinks to correct GDS lib and include directories so there
is no need to hardcode full path that is architecture-dependent. Set
default GDS path to point to CUDA install directory and use 'include' and
'lib64' symlinks.

Signed-off-by: Vlad Buslov <vladbu@nvidia.com>

* Fix container build to properly support x86 and Arm

Currently the build script and Docker build file assume that the container
is built on x86_64 Linux host (by not specifying a platform which makes
docker default to host architecture, hardcoding 'x86_64' in several places,
etc.), which makes it impossible to create either x86_64 or aarch64
container on Arm host or aarch64 container on x86 host.

Configure target architecture via 'ARCH' docker variable. Set it to x86 by
default in Dockerfile (for any users that use it directly) and to host
architecture in build-container.sh. Allow user to specify ARCH value via
'--arch' CLI parameter, set docker build platform value accordingly and
pass the value to docker build as a build arg.

Signed-off-by: Vlad Buslov <vladbu@nvidia.com>

* Fix nixlbench container build to properly support x86 and Arm

Similar to nixl library container infrastructure nixlbench also does not
specify a docker build platform, hardcodes 'x86_64' in several places and
assumes x86 manylinux platform which makes it impossible to build anything
on aarch64 host or for aarch64 target.

Configure target architecture via 'ARCH' docker variable. Set it to x86 by
default in nixlbench Dockerfile (for any users that use it directly) and to
host architecture in build-container.sh. Allow user to specify ARCH value
via '--arch' CLI parameter, set docker build platform value accordingly and
pass the value to docker build as a build arg.

Signed-off-by: Vlad Buslov <vladbu@nvidia.com>

---------

Signed-off-by: Vlad Buslov <vladbu@nvidia.com>

Author: w1ldptr

benchmark/nixlbench/contrib/Dockerfile

@@ -73,9 +73,10 @@ FROM ucx_${UCX}_image AS ucx_image
 
 # --- Stage 4: Final Image Assembly ---
 # Re-declare ARGs needed in this final stage
+ARG ARCH="x86_64"
 ARG DEFAULT_PYTHON_VERSION
 ARG WHL_PYTHON_VERSIONS="3.12"
-ARG WHL_PLATFORM="manylinux_2_39_x86_64"
+ARG WHL_PLATFORM="manylinux_2_39_$ARCH"
 
 WORKDIR /workspace
 RUN git clone https://github.com/etcd-cpp-apiv3/etcd-cpp-apiv3.git &&\
@@ -104,9 +105,9 @@ RUN rm -rf build && \
     ninja && \
     ninja install
 
-ENV NIXL_PLUGIN_DIR=/usr/local/nixl/lib/x86_64-linux-gnu/plugins
-RUN echo "/usr/local/nixl/lib/x86_64-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \
-    echo "/usr/local/nixl/lib/x86_64-linux-gnu/plugins" >> /etc/ld.so.conf.d/nixl.conf && \
+ENV NIXL_PLUGIN_DIR=/usr/local/nixl/lib/$ARCH-linux-gnu/plugins
+RUN echo "/usr/local/nixl/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \
+    echo "/usr/local/nixl/lib/$ARCH-linux-gnu/plugins" >> /etc/ld.so.conf.d/nixl.conf && \
     ldconfig
 
 # Create the wheel

benchmark/nixlbench/contrib/build.sh

@@ -35,7 +35,10 @@ fi
 
 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base
 BASE_IMAGE_TAG=25.03-cuda12.8-devel-ubuntu24.04
-WHL_PLATFORM=manylinux_2_39_x86_64
+ARCH=$(uname -m)
+[ "$ARCH" = "arm64" ] && ARCH="aarch64"
+WHL_BASE=manylinux_2_39
+WHL_PLATFORM=${WHL_BASE}_${ARCH}
 WHL_PYTHON_VERSIONS="3.12"
 OS="ubuntu24"
 
@@ -108,6 +111,15 @@ get_options() {
                 missing_requirement $1
             fi
             ;;
+        --arch)
+            if [ "$2" ]; then
+                ARCH=$2
+                WHL_PLATFORM=${WHL_BASE}_${ARCH}
+                shift
+            else
+                missing_requirement $1
+            fi
+            ;;
         --)
             shift
             break
@@ -147,6 +159,7 @@ show_build_options() {
     echo "Build Context: ${BUILD_CONTEXT}"
     echo "Build Context Args: ${BUILD_CONTEXT_ARGS}"
     echo "Base Image: ${BASE_IMAGE}:${BASE_IMAGE_TAG}"
+    echo "Container arch: ${ARCH}"
     echo "Python Versions for wheel build: ${WHL_PYTHON_VERSIONS}"
     echo "Wheel Platform: ${WHL_PLATFORM}"
 }
@@ -161,6 +174,7 @@ show_help() {
     echo "  [--os [ubuntu24|ubuntu22] to select Ubuntu version]"
     echo "  [--python-versions python versions to build for, comma separated]"
     echo "  [--tag tag for image]"
+    echo "  [--arch [x86_64|aarch64] to select target architecture]"
     exit 0
 }
 
@@ -178,7 +192,8 @@ get_options "$@"
 BUILD_ARGS+=" --build-arg BASE_IMAGE=$BASE_IMAGE --build-arg BASE_IMAGE_TAG=$BASE_IMAGE_TAG"
 BUILD_ARGS+=" --build-arg WHL_PYTHON_VERSIONS=$WHL_PYTHON_VERSIONS"
 BUILD_ARGS+=" --build-arg WHL_PLATFORM=$WHL_PLATFORM"
+BUILD_ARGS+=" --build-arg ARCH=$ARCH"
 
 show_build_options
 
-docker build -f $DOCKER_FILE $BUILD_ARGS $TAG $NO_CACHE $BUILD_ARGS $BUILD_CONTEXT_ARGS $BUILD_CONTEXT --progress plain
+docker build --platform linux/$ARCH -f $DOCKER_FILE $BUILD_ARGS $TAG $NO_CACHE $BUILD_CONTEXT_ARGS $BUILD_CONTEXT --progress plain

contrib/Dockerfile

@@ -18,6 +18,7 @@ ARG BASE_IMAGE_TAG="25.03-cuda12.8-devel-ubuntu24.04"
 
 FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG}
 
+ARG ARCH="x86_64"
 ARG DEFAULT_PYTHON_VERSION="3.12"
 
 RUN apt-get update -y && \
@@ -57,13 +58,16 @@ ENV RUSTUP_HOME=/usr/local/rustup \
     CARGO_HOME=/usr/local/cargo \
     PATH=/usr/local/cargo/bin:$PATH \
     RUST_VERSION=1.86.0 \
-    RUSTARCH=x86_64-unknown-linux-gnu
+    RUSTARCH=${ARCH}-unknown-linux-gnu
 
-RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
-    echo "a3339fb004c3d0bb9862ba0bce001861fe5cbde9c10d16591eb3f39ee6cd3e7f *rustup-init" | sha256sum -c - && \
+# Download rustup-init and its checksum for the target architecture
+RUN wget --tries=3 --waitretry=5 \
+    "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" \
+    "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init.sha256" && \
+    sha256sum -c rustup-init.sha256 && \
     chmod +x rustup-init && \
     ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
-    rm rustup-init && \
+    rm rustup-init* && \
     chmod -R a+w $RUSTUP_HOME $CARGO_HOME
 
 WORKDIR /workspace/nixl
@@ -110,16 +114,16 @@ RUN rm -rf build && \
     ninja install
 
 ENV NIXL_PREFIX=/usr/local/nixl
-ENV NIXL_PLUGIN_DIR=/usr/local/nixl/lib/x86_64-linux-gnu/plugins
-RUN echo "/usr/local/nixl/lib/x86_64-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \
-    echo "/usr/local/nixl/lib/x86_64-linux-gnu/plugins" >> /etc/ld.so.conf.d/nixl.conf && \
+ENV NIXL_PLUGIN_DIR=/usr/local/nixl/lib/$ARCH-linux-gnu/plugins
+RUN echo "/usr/local/nixl/lib/$ARCH-linux-gnu" > /etc/ld.so.conf.d/nixl.conf && \
+    echo "/usr/local/nixl/lib/$ARCH-linux-gnu/plugins" >> /etc/ld.so.conf.d/nixl.conf && \
     ldconfig
 
 RUN cd src/bindings/rust && cargo build --release --locked
 
 # Create the wheel
 ARG WHL_PYTHON_VERSIONS="3.12"
-ARG WHL_PLATFORM="manylinux_2_39_x86_64"
+ARG WHL_PLATFORM="manylinux_2_39_$ARCH"
 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; \

contrib/build-container.sh

@@ -30,7 +30,10 @@ VERSION=v$latest_tag.dev.$commit_id
 
 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base
 BASE_IMAGE_TAG=25.03-cuda12.8-devel-ubuntu24.04
-WHL_PLATFORM=manylinux_2_39_x86_64
+ARCH=$(uname -m)
+[ "$ARCH" = "arm64" ] && ARCH="aarch64"
+WHL_BASE=manylinux_2_39
+WHL_PLATFORM=${WHL_BASE}_${ARCH}
 WHL_PYTHON_VERSIONS="3.12"
 OS="ubuntu24"
 
@@ -84,6 +87,15 @@ get_options() {
                 missing_requirement $1
             fi
             ;;
+        --arch)
+            if [ "$2" ]; then
+                ARCH=$2
+                WHL_PLATFORM=${WHL_BASE}_${ARCH}
+                shift
+            else
+                missing_requirement $1
+            fi
+            ;;
         --)
             shift
             break
@@ -103,7 +115,7 @@ get_options() {
 
     if [[ $OS == "ubuntu22" ]]; then
         BASE_IMAGE_TAG=24.10-cuda12.6-devel-ubuntu22.04
-        WHL_PLATFORM=manylinux_2_34_x86_64
+        WHL_PLATFORM=manylinux_2_34_${ARCH}
     fi
 
     if [ -z "$TAG" ]; then
@@ -117,6 +129,7 @@ show_build_options() {
     echo "Image Tag: ${TAG}"
     echo "Build Context: ${BUILD_CONTEXT}"
     echo "Base Image: ${BASE_IMAGE}:${BASE_IMAGE_TAG}"
+    echo "Container arch: ${ARCH}"
     echo "Python Versions for wheel build: ${WHL_PYTHON_VERSIONS}"
     echo "Wheel Platform: ${WHL_PLATFORM}"
 }
@@ -129,6 +142,7 @@ show_help() {
     echo "  [--os [ubuntu24|ubuntu22] to select Ubuntu version]"
     echo "  [--tag tag for image]"
     echo "  [--python-versions python versions to build for, comma separated]"
+    echo "  [--arch [x86_64|aarch64] to select target architecture]"
     exit 0
 }
 
@@ -151,7 +165,8 @@ fi
 BUILD_ARGS+=" --build-arg BASE_IMAGE=$BASE_IMAGE --build-arg BASE_IMAGE_TAG=$BASE_IMAGE_TAG"
 BUILD_ARGS+=" --build-arg WHL_PYTHON_VERSIONS=$WHL_PYTHON_VERSIONS"
 BUILD_ARGS+=" --build-arg WHL_PLATFORM=$WHL_PLATFORM"
+BUILD_ARGS+=" --build-arg ARCH=$ARCH"
 
 show_build_options
 
-docker build -f $DOCKER_FILE $BUILD_ARGS $TAG $NO_CACHE $BUILD_ARGS $BUILD_CONTEXT
+docker build --platform linux/$ARCH -f $DOCKER_FILE $BUILD_ARGS $TAG $NO_CACHE $BUILD_CONTEXT

meson_options.txt

@@ -19,7 +19,7 @@ option('etcd_lib_path', type: 'string', value: '', description: 'Path to ETCD Li
 option('disable_gds_backend', type : 'boolean', value : false, description : 'disable gds backend')
 option('disable_mooncake_backend', type : 'boolean', value : false, description : 'disable mooncake backend')
 option('install_headers', type : 'boolean', value : true, description : 'install headers')
-option('gds_path', type: 'string', value: '/usr/local/cuda/targets/x86_64-linux/', description: 'Path to GDS CuFile install')
+option('gds_path', type: 'string', value: '/usr/local/cuda/', description: 'Path to GDS CuFile install')
 option('cudapath_inc', type: 'string', value: '', description: 'Include path for CUDA')
 option('cudapath_lib', type: 'string', value: '', description: 'Library path for CUDA')
 option('cudapath_stub', type: 'string', value: '', description: 'Extra Stub path for CUDA')

src/plugins/cuda_gds/meson.build

@@ -15,7 +15,7 @@
 
 gds_path = get_option('gds_path')
 if gds_path != ''
-    gds_lib_path = gds_path + '/lib'
+    gds_lib_path = gds_path + '/lib64'
     gds_inc_path = gds_path + '/include'
     cufile_dep = declare_dependency(
         link_args: ['-L' + gds_lib_path, '-lcufile'],