huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 19:34:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(dockerfile): fix nvidia env vars

Browse Source
This commit is contained in:
OlivierDehaene 2023-05-09 11:35:06 +02:00
parent b4aa87db58
commit 5ef4c94289
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
Dockerfile
View File

@ -120,6 +120,16 @@ ENV HUGGINGFACE_HUB_CACHE=/data \
NUM_SHARD=1 \
PORT=80
# NVIDIA env vars
ENV NVIDIA_REQUIRE_CUDA cuda>=11.8
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
# Required for nvidia-docker v1
RUN /bin/bash -c echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf && \
echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
LABEL com.nvidia.volumes.needed="nvidia_driver"
WORKDIR /usr/src
@ -165,11 +175,6 @@ FROM base as sagemaker
COPY sagemaker-entrypoint.sh entrypoint.sh
RUN chmod +x entrypoint.sh
# NVIDIA env vars
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
ENTRYPOINT ["./entrypoint.sh"]
# Final image