Put nvidia-smi on path for cuda image variants (#2124)

* Put nvidia-smi on path for cuda image variants * Improve comments style --------- Co-authored-by: Ayaz Salikhov <mathbunnyru@users.noreply.github.com>
jupyter · Jul 1, 2024 · 3fef154 · 3fef154
1 parent 37c03eb
commit 3fef154
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 2 deletions.
diff --git a/images/pytorch-notebook/cuda11/Dockerfile b/images/pytorch-notebook/cuda11/Dockerfile
@@ -23,3 +23,8 @@ RUN pip install --no-cache-dir --extra-index-url=https://pypi.nvidia.com --index
 # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html#dockerfiles
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+# Puts the nvidia-smi binary (system management interface) on path
+# with associated library files to execute it
+ENV PATH=${PATH}:/usr/local/nvidia/bin
+ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/nvidia/lib64
diff --git a/images/pytorch-notebook/cuda12/Dockerfile b/images/pytorch-notebook/cuda12/Dockerfile
@@ -23,3 +23,8 @@ RUN pip install --no-cache-dir --extra-index-url=https://pypi.nvidia.com --index
 # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html#dockerfiles
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+# Puts the nvidia-smi binary (system management interface) on path
+# with associated library files to execute it
+ENV PATH=${PATH}:/usr/local/nvidia/bin
+ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/nvidia/lib64
diff --git a/images/tensorflow-notebook/cuda/Dockerfile b/images/tensorflow-notebook/cuda/Dockerfile
@@ -25,3 +25,8 @@ COPY --chown="${NB_UID}:${NB_GID}" nvidia-lib-dirs.sh "${CONDA_DIR}/etc/conda/ac
 # https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html#dockerfiles
 ENV NVIDIA_VISIBLE_DEVICES="all" \
     NVIDIA_DRIVER_CAPABILITIES="compute,utility"
+
+# Puts the nvidia-smi binary (system management interface) on path
+# with associated library files to execute it
+ENV PATH=${PATH}:/usr/local/nvidia/bin
+ENV LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/nvidia/lib64
diff --git a/images/tensorflow-notebook/cuda/nvidia-lib-dirs.sh b/images/tensorflow-notebook/cuda/nvidia-lib-dirs.sh
@@ -2,8 +2,8 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
-# This adds the NVIDIA libraries to the LD_LIBRARY_PATH. Workaround for
-# https://github.com/tensorflow/tensorflow/issues/63362
+# This adds NVIDIA Python package libraries to the LD_LIBRARY_PATH.
+# Workaround for https://github.com/tensorflow/tensorflow/issues/63362
 NVIDIA_DIR=$(dirname "$(python -c 'import nvidia;print(nvidia.__file__)')")
 LD_LIBRARY_PATH=$(echo "${NVIDIA_DIR}"/*/lib/ | sed -r 's/\s+/:/g')${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
 export LD_LIBRARY_PATH