Skip to content

Build and container fixes for aarch64 #408

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 2, 2025
Merged

Build and container fixes for aarch64 #408

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
48d8093
Fix gds_path default
w1ldptr May 30, 2025
b5f512e
Fix container build to properly support x86 and Arm
w1ldptr May 30, 2025
b631cd8
Fix nixlbench container build to properly support x86 and Arm
w1ldptr May 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions benchmark/nixlbench/contrib/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 &&\
Expand Down Expand Up @@ -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
Expand Down
19 changes: 17 additions & 2 deletions benchmark/nixlbench/contrib/build.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}"
}
Expand All @@ -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
}

Expand All @@ -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
20 changes: 12 additions & 8 deletions contrib/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 && \
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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; \
Expand Down
21 changes: 18 additions & 3 deletions contrib/build-container.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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}"
}
Expand All @@ -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
}

Expand All @@ -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
2 changes: 1 addition & 1 deletion meson_options.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
2 changes: 1 addition & 1 deletion src/plugins/cuda_gds/meson.build
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'],
Expand Down