makaveli10
commited on
Commit
•
25adfbc
1
Parent(s):
8accb9d
add cuda-arch as a build-arg
Browse files- docker/base-image/Dockerfile +2 -1
- docker/base-image/install-trt-llm.sh +3 -1
- docker/build.sh +8 -1
docker/base-image/Dockerfile
CHANGED
@@ -4,9 +4,10 @@ ARG BASE_IMAGE=nvcr.io/nvidia/cuda
|
|
4 |
ARG BASE_TAG=12.2.2-devel-ubuntu22.04
|
5 |
|
6 |
FROM ${BASE_IMAGE}:${BASE_TAG} as base
|
7 |
-
|
8 |
WORKDIR /root
|
9 |
COPY install-deps.sh /root
|
|
|
10 |
RUN bash install-deps.sh && rm install-deps.sh
|
11 |
|
12 |
COPY install-trt-llm.sh /root
|
|
|
4 |
ARG BASE_TAG=12.2.2-devel-ubuntu22.04
|
5 |
|
6 |
FROM ${BASE_IMAGE}:${BASE_TAG} as base
|
7 |
+
ARG CUDA_ARCH
|
8 |
WORKDIR /root
|
9 |
COPY install-deps.sh /root
|
10 |
+
ENV CUDA_ARCH=${CUDA_ARCH}
|
11 |
RUN bash install-deps.sh && rm install-deps.sh
|
12 |
|
13 |
COPY install-trt-llm.sh /root
|
docker/base-image/install-trt-llm.sh
CHANGED
@@ -3,8 +3,10 @@
|
|
3 |
export ENV=${ENV:-/etc/shinit_v2}
|
4 |
source $ENV
|
5 |
|
|
|
|
|
6 |
cd /root/TensorRT-LLM
|
7 |
-
python3 scripts/build_wheel.py --clean --cuda_architectures "
|
8 |
pip install build/tensorrt_llm-0.7.1-cp310-cp310-linux_x86_64.whl
|
9 |
mv examples ../TensorRT-LLM-examples
|
10 |
cd ..
|
|
|
3 |
export ENV=${ENV:-/etc/shinit_v2}
|
4 |
source $ENV
|
5 |
|
6 |
+
CUDA_ARCH="${CUDA_ARCH:-89-real;90-real}"
|
7 |
+
|
8 |
cd /root/TensorRT-LLM
|
9 |
+
python3 scripts/build_wheel.py --clean --cuda_architectures "$CUDA_ARCH" --trt_root /usr/local/tensorrt
|
10 |
pip install build/tensorrt_llm-0.7.1-cp310-cp310-linux_x86_64.whl
|
11 |
mv examples ../TensorRT-LLM-examples
|
12 |
cd ..
|
docker/build.sh
CHANGED
@@ -1,10 +1,17 @@
|
|
1 |
#!/bin/bash -e
|
2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3 |
[ -n "$VERBOSE" ] && ARGS="--progress plain"
|
4 |
|
5 |
(
|
6 |
cd base-image &&
|
7 |
-
docker build $ARGS -t ghcr.io/collabora/whisperfusion-base:latest .
|
8 |
)
|
9 |
|
10 |
mkdir -p scratch-space
|
|
|
1 |
#!/bin/bash -e
|
2 |
|
3 |
+
if [ -n "$1" ]; then
|
4 |
+
CUDA_ARCH="$1"
|
5 |
+
BASE_IMAGE_BUILD_ARG="--build-arg CUDA_ARCH=$CUDA_ARCH"
|
6 |
+
else
|
7 |
+
BASE_IMAGE_BUILD_ARG=""
|
8 |
+
fi
|
9 |
+
|
10 |
[ -n "$VERBOSE" ] && ARGS="--progress plain"
|
11 |
|
12 |
(
|
13 |
cd base-image &&
|
14 |
+
docker build $ARGS $BASE_IMAGE_BUILD_ARG -t ghcr.io/collabora/whisperfusion-base:latest .
|
15 |
)
|
16 |
|
17 |
mkdir -p scratch-space
|