feat: Add unified x86 / aarch64 (ARM) build for TRTLLM image (#803)

Signed-off-by: Ryan McCormick <rmccormick@nvidia.com>

container/Dockerfile.tensorrt_llm

@@ -15,25 +15,45 @@
 
 ARG BASE_IMAGE="tensorrt_llm/release"
 ARG BASE_IMAGE_TAG="latest_squashed"
-ARG MANYLINUX_IMAGE="quay.io/pypa/manylinux_2_28_x86_64"
 ARG RELEASE_BUILD
 
+# Define general architecture ARGs for supporting both x86 and aarch64 builds.
+#   ARCH: Used for package suffixes (e.g., amd64, arm64)
+#   ARCH_ALT: Used for Rust targets, manylinux suffix (e.g., x86_64, aarch64)
+#
+# Default values are for x86/amd64:
+#   --build-arg ARCH=amd64 --build-arg ARCH_ALT=x86_64
+#
+# For arm64/aarch64, build with:
+#   --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64
+#
+# NOTE: There isn't an easy way to define one of these values based on the other value
+# without adding if statements everywhere, so just define both as ARGs for now.
+ARG ARCH=amd64
+ARG ARCH_ALT=x86_64
+
 ##################################
 ########## Build Image ###########
 ##################################
 
 FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS build
 
+# Redeclare ARCH and ARCH_ALT so they're available in this build stage
+ARG ARCH
+ARG ARCH_ALT
+
 USER root
 
 # Install utilities
 RUN apt update -y && apt install -y git wget curl nvtop tmux vim
+
 # nats
-RUN wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/v2.10.24/nats-server-v2.10.24-amd64.deb && \
-    dpkg -i nats-server-v2.10.24-amd64.deb && rm nats-server-v2.10.24-amd64.deb
+RUN wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/v2.10.24/nats-server-v2.10.24-${ARCH}.deb && \
+    dpkg -i nats-server-v2.10.24-${ARCH}.deb && rm nats-server-v2.10.24-${ARCH}.deb
+
 # etcd
 ENV ETCD_VERSION="v3.5.18"
-RUN wget https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-amd64.tar.gz -O /tmp/etcd.tar.gz && \
+RUN wget https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
     mkdir -p /usr/local/bin/etcd && \
     tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
     rm /tmp/etcd.tar.gz
@@ -75,11 +95,14 @@ RUN apt-get update && \
 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
+    RUST_VERSION=1.86.0
 
+# Define Rust target based on ARCH_ALT ARG
+ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu
+
+# Install Rust using RUSTARCH derived from ARCH_ALT
 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 - && \
+    # TODO: Add SHA check back based on RUSTARCH
     chmod +x rustup-init && \
     ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
     rm rustup-init && \
@@ -104,8 +127,10 @@ RUN mkdir /opt/dynamo && \
 ####### WHEEL BUILD STAGE #########
 ###################################
 
-# Build the wheel in the manylinux environment
-FROM ${MANYLINUX_IMAGE} AS wheel_builder
+# Redeclare ARCH_ALT ARG so it's available for interpolation in the FROM instruction
+ARG ARCH_ALT
+
+FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
 ARG RELEASE_BUILD
 ARG CARGO_BUILD_JOBS
 # Set CARGO_BUILD_JOBS to 16 if not provided
@@ -118,8 +143,6 @@ WORKDIR /workspace
 RUN yum update -y \
     && yum install -y python3.12-devel \
     && yum install -y protobuf-compiler \
-    || yum install -y https://raw.repo.almalinux.org/almalinux/8.10/AppStream/x86_64/os/Packages/protobuf-3.5.0-15.el8.x86_64.rpm \
-    https://raw.repo.almalinux.org/almalinux/8.10/AppStream/x86_64/os/Packages/protobuf-compiler-3.5.0-15.el8.x86_64.rpm \
     && yum clean all \
     && rm -rf /var/cache/yum
 
@@ -157,7 +180,7 @@ RUN cargo build --release --locked --features mistralrs,sglang,vllm,python \
     -p file://$PWD/components/http \
     -p metrics
 
-# Build dynamo wheel
+# Build dynamo wheels
 RUN uv build --wheel --out-dir /workspace/dist && \
     cd /workspace/lib/bindings/python && \
     uv build --wheel --out-dir /workspace/dist --python 3.12 && \
@@ -200,12 +223,12 @@ RUN mkdir -p /opt/dynamo/bindings/wheels && \
     cp target/release/metrics /usr/local/bin && \
     cp target/release/mock_worker /usr/local/bin
 
+# Install wheels
 RUN . /opt/dynamo/venv/bin/activate && \
     uv pip install /workspace/dist/ai_dynamo_runtime*cp312*.whl && \
     uv pip install /workspace/dist/ai_dynamo*any.whl
 
-# Install dynamo.runtime and dynamo.llm wheels globally in container for tests that
-# currently run without virtual environment activated.
+# Install dynamo.runtime and dynamo.llm wheels globally in container for tests
 # TODO: In future, we may use a virtualenv for everything and remove this.
 RUN pip install dist/ai_dynamo_runtime*cp312*.whl  && \
     pip install dist/ai_dynamo*any.whl

container/build.sh

@@ -283,6 +283,11 @@ error() {
 
 get_options "$@"
 
+# Automatically set ARCH and ARCH_ALT if PLATFORM is linux/arm64
+if [[ "$PLATFORM" == *"linux/arm64"* ]]; then
+    BUILD_ARGS+=" --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64 "
+fi
+
 # Update DOCKERFILE if framework is VLLM
 if [[ $FRAMEWORK == "VLLM" ]]; then
     DOCKERFILE=${SOURCE_DIR}/Dockerfile.vllm