From 75001ae6b7a8dcc19ae9a125aa156e30a9609427 Mon Sep 17 00:00:00 2001
From: makaveli10 <vineet.suryan@collabora.com>
Date: Fri, 19 Jan 2024 11:58:40 +0000
Subject: [PATCH] updatetensorrt-llm dockerfile

---
 docker/Dockerfile.tensorrt | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt
index 9bcf7b0..c91dc8e 100644
--- a/docker/Dockerfile.tensorrt
+++ b/docker/Dockerfile.tensorrt
@@ -1,8 +1,14 @@
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
 ARG BASE_TAG=23.10-py3
+ARG CUDA_ARCHITECTURES="89-real"
 
 FROM ${BASE_IMAGE}:${BASE_TAG} as base
 
 WORKDIR /home
 COPY script/install_trt_llm.sh install_trt_llm.sh
-RUN bash install_trt_llm.sh && rm install_trt_llm.sh
\ No newline at end of file
+RUN bash install_trt_llm.sh ${CUDA_ARCHITECTURES} && rm install_trt_llm.sh
+
+COPY scripts/setup.sh /home
+RUN apt update && bash setup.sh && rm setup.sh
+COPY requirements/ /home
+RUN pip install -r server.txt && rm server.txt