add cuda-arch as a build-arg

docker/base-image/Dockerfile

@@ -4,9 +4,10 @@ ARG BASE_IMAGE=nvcr.io/nvidia/cuda
 ARG BASE_TAG=12.2.2-devel-ubuntu22.04
 
 FROM ${BASE_IMAGE}:${BASE_TAG} as base
-
+ARG CUDA_ARCH
 WORKDIR /root
 COPY install-deps.sh /root
+ENV CUDA_ARCH=${CUDA_ARCH}
 RUN bash install-deps.sh && rm install-deps.sh
 
 COPY install-trt-llm.sh /root

docker/base-image/install-trt-llm.sh

@@ -3,8 +3,10 @@
 export ENV=${ENV:-/etc/shinit_v2}
 source $ENV
 
+CUDA_ARCH="${CUDA_ARCH:-89-real;90-real}"
+
 cd /root/TensorRT-LLM
-python3 scripts/build_wheel.py --clean --cuda_architectures "89-real;90-real" --trt_root /usr/local/tensorrt
+python3 scripts/build_wheel.py --clean --cuda_architectures "$CUDA_ARCH" --trt_root /usr/local/tensorrt
 pip install build/tensorrt_llm-0.7.1-cp310-cp310-linux_x86_64.whl
 mv examples ../TensorRT-LLM-examples
 cd ..

docker/build.sh

@@ -1,10 +1,17 @@
 #!/bin/bash -e
 
+if [ -n "$1" ]; then
+  CUDA_ARCH="$1"
+  BASE_IMAGE_BUILD_ARG="--build-arg CUDA_ARCH=$CUDA_ARCH"
+else
+  BASE_IMAGE_BUILD_ARG=""
+fi
+
 [ -n "$VERBOSE" ] && ARGS="--progress plain"
 
 (
   cd base-image &&
-  docker build $ARGS -t ghcr.io/collabora/whisperfusion-base:latest .
+  docker build $ARGS $BASE_IMAGE_BUILD_ARG -t ghcr.io/collabora/whisperfusion-base:latest .
 )
 
 mkdir -p scratch-space