fix(dockerfile): fix nvidia env vars (#297)

Fixes #291
This commit is contained in:
OlivierDehaene 2023-05-09 12:36:02 +02:00 committed by GitHub
parent b4aa87db58
commit e9b01b3433
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 5 deletions

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