ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
ARG BASE_TAG=23.10-py3

FROM ${BASE_IMAGE}:${BASE_TAG} as base

WORKDIR /home
COPY scripts/install_tensorrt_llm.sh install_tensorrt_llm.sh
RUN bash install_tensorrt_llm.sh && rm install_tensorrt_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