Move PyTorch to Conda to prevent collision

de8852a2 · Seungmin Kim · d7cc9251 · d7cc9251 · de8852a2 · de8852a2
Commit de8852a2 authored 2 months ago by Seungmin Kim
--- a/20tensorboard-proxy-env.sh
+++ b/20tensorboard-proxy-env.sh
-#!/bin/bash
-# Copyright (c) Jupyter Development Team.
-# Distributed under the terms of the Modified BSD License.
-
-set -e
-
-# Initialize the TENSORBOARD_PROXY_URL with the appropriate path
-# to use jupyter-server-proxy.
-
-export TENSORBOARD_PROXY_URL="${JUPYTERHUB_SERVICE_PREFIX:-/}proxy/%PORT%/"
--- a/Dockerfile
+++ b/Dockerfile
-ARG BASE_CONTAINER=quay.io/jupyter/scipy-notebook:2024-08-26
+ARG BASE_CONTAINER=quay.io/jupyter/tensorflow-notebook:2025-01-27
 FROM $BASE_CONTAINER

 LABEL maintainer="Dima Mishin <dmishin@ucsd.edu>, Seungmin Kim <ehf@yonsei.ac.kr>"

 SHELL ["/bin/bash", "-o", "pipefail", "-c"]

-USER root
-
-# Install code-server
-RUN apt-get -q update \
-    && apt-get install -yq --no-install-recommends \
-        curl \
-        jq \
-    && CODESERVER_VERSION="$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | jq -r '.tag_name' | sed 's/[^0-9\.\-]*//g')" \
-    && curl -o code-server.deb -fsSL "https://github.com/coder/code-server/releases/download/v${CODESERVER_VERSION}/code-server_${CODESERVER_VERSION}_$(dpkg --print-architecture).deb" \
-    && apt-get install -yq --no-install-recommends ./code-server.deb && rm -f ./code-server.deb \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-
 USER $NB_UID

 RUN mamba install -c conda-forge --yes \
        'curl' \
        'ghostscript' \
        'graphviz' \
+        'jq' \
        'rasterio' \
        'gdal' \
        'libgdal-arrow-parquet' \
+    mamba install -c pytorch --yes \
+        'pytorch' \
+        'torchvision' \
+        'torchaudio' \
+        'cpuonly' \
    && mamba clean --all -f -y \
    && if [ "$(uname -m)" = "x86_64" ]; then TF_POSTFIX="-cpu"; else TF_POSTFIX=""; fi \
-    && pip install --no-cache-dir --extra-index-url "https://download.pytorch.org/whl/cpu" \
-        # PyTorch with CPU
-        torch \
-        torchvision \
-        torchaudio \
-        # Tensorflow with CPU
-        "jupyter-server-proxy" \
-        "tensorflow${TF_POSTFIX}" \
-        # JAX and Flax with CPU
-        jax \
-        flax \
+    && pip install --no-cache-dir \
+        "tensorflow${TF_POSTFIX}==$(python3 -c 'import tensorflow as tf; print(tf.__version__)' 2>/dev/null)" \
        # Keras and other packages
        keras \
        keras-hub \
        keras-nlp \
        keras-tuner \
+        # JAX and Flax with CPU
+        jax \
+        flax \
        # https://github.com/explosion/cython-blis/issues/117#issuecomment-2599094860
        'spacy<3.8.0' \
        accelerate \
@@ -69,6 +57,15 @@ RUN mamba install -c conda-forge --yes \
    && fix-permissions "${CONDA_DIR}" \
    && fix-permissions "/home/${NB_USER}"

-COPY --chown="${NB_UID}:${NB_GID}" 20tensorboard-proxy-env.sh /usr/local/bin/before-notebook.d/
+USER root
+
+# Install code-server
+RUN CODESERVER_VERSION="$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | jq -r '.tag_name' | sed 's/[^0-9\.\-]*//g')" \
+    && curl -o code-server.deb -fsSL "https://github.com/coder/code-server/releases/download/v${CODESERVER_VERSION}/code-server_${CODESERVER_VERSION}_$(dpkg --print-architecture).deb" \
+    && apt-get -q update \
+    && apt-get install -yq --no-install-recommends ./code-server.deb && rm -f ./code-server.deb \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+USER $NB_UID

 WORKDIR $HOME
--- a/Dockerfile_cuda
+++ b/Dockerfile_cuda
-ARG BASE_CONTAINER=quay.io/jupyter/scipy-notebook:2024-08-26
+ARG BASE_CONTAINER=quay.io/jupyter/tensorflow-notebook:cuda-2025-01-27
 FROM $BASE_CONTAINER

 LABEL maintainer="Dima Mishin <dmishin@ucsd.edu>, Seungmin Kim <ehf@yonsei.ac.kr>"

 SHELL ["/bin/bash", "-o", "pipefail", "-c"]

-USER root
-
-# Install CUDA, https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
-ARG CUDA_VERSION=12.4
-ARG CUDA_TENSORRT_VERSION=12.4
-RUN apt-get -q update && \
-    apt-get install -yq --no-install-recommends \
-    gnupg curl ca-certificates && \
-    # Install key and repository
-    curl -fsSL -O "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu$(grep '^VERSION_ID'= /etc/os-release | cut -d= -f2 | tr -d '\"\.')/$(dpkg --print-architecture | sed -e 's/arm64/sbsa/'  -e 's/i.*86/i386/' -e 's/amd64/x86_64/' -e 's/unknown/x86_64/')/cuda-keyring_1.1-1_all.deb" && \
-    apt-get install -yq --no-install-recommends \
-    ./cuda-keyring_1.1-1_all.deb && \
-    rm -f ./cuda-keyring_1.1-1_all.deb && \
-    # Install CUDA packages
-    apt-get -q update && \
-    apt-get install -yq --no-install-recommends \
-    cuda-cudart-${CUDA_VERSION//./-} \
-    cuda-cudart-dev-${CUDA_VERSION//./-} \
-    cuda-command-line-tools-${CUDA_VERSION//./-} \
-    cuda-libraries-${CUDA_VERSION//./-} \
-    cuda-libraries-dev-${CUDA_VERSION//./-} \
-    cuda-minimal-build-${CUDA_VERSION//./-} \
-    cuda-nsight-compute-${CUDA_VERSION//./-} \
-    cuda-nvml-dev-${CUDA_VERSION//./-} \
-    cuda-nvprof-${CUDA_VERSION//./-} \
-    libnpp-${CUDA_VERSION//./-} \
-    libnpp-dev-${CUDA_VERSION//./-} \
-    cuda-nvtx-${CUDA_VERSION//./-} \
-    libcusparse-${CUDA_VERSION//./-} \
-    libcusparse-dev-${CUDA_VERSION//./-} \
-    libcublas-${CUDA_VERSION//./-} \
-    libcublas-dev-${CUDA_VERSION//./-} \
-    libnccl2=*+cuda${CUDA_VERSION} \
-    libnccl-dev=*+cuda${CUDA_VERSION} \
-    # Install CUDNN packages
-    libcudnn9-cuda-${CUDA_VERSION%%.*} \
-    libcudnn9-dev-cuda-${CUDA_VERSION%%.*} \
-    # Install TensorRT backend
-    python3-libnvinfer=*+cuda${CUDA_TENSORRT_VERSION} \
-    python3-libnvinfer-lean=*+cuda${CUDA_TENSORRT_VERSION} \
-    python3-libnvinfer-dispatch=*+cuda${CUDA_TENSORRT_VERSION} \
-    python3-libnvinfer-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-lean10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-lean-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-dispatch10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-dispatch-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-headers-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-headers-plugin-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-plugin10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-plugin-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-vc-plugin10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvinfer-vc-plugin-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvonnxparsers10=*+cuda${CUDA_TENSORRT_VERSION} \
-    libnvonnxparsers-dev=*+cuda${CUDA_TENSORRT_VERSION} \
-    ocl-icd-libopencl1 && \
-    apt-mark hold libcublas-${CUDA_VERSION//./-} libcublas-dev-${CUDA_VERSION//./-} libnccl2 libnccl-dev libcudnn9-cuda-${CUDA_VERSION%%.*} libcudnn9-dev-cuda-${CUDA_VERSION%%.*} python3-libnvinfer-dev && \
-    apt-get clean && rm -rf /var/lib/apt/lists/* && \
-    mkdir -pm755 /etc/OpenCL/vendors && echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \
-    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 PATH="/usr/local/nvidia/bin:/usr/local/cuda/bin${PATH:+:${PATH}}"
-ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}/usr/local/nvidia/lib:/usr/local/nvidia/lib64"
-ENV LIBRARY_PATH="${LIBRARY_PATH:+:${LIBRARY_PATH}}:/usr/local/cuda/lib64/stubs"
-
-ENV NVIDIA_VISIBLE_DEVICES=all
-ENV NVIDIA_DRIVER_CAPABILITIES=all
-
-# Install code-server
-RUN apt-get -q update \
-    && apt-get install -yq --no-install-recommends \
-        curl \
-        jq \
-    && CODESERVER_VERSION="$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | jq -r '.tag_name' | sed 's/[^0-9\.\-]*//g')" \
-    && curl -o code-server.deb -fsSL "https://github.com/coder/code-server/releases/download/v${CODESERVER_VERSION}/code-server_${CODESERVER_VERSION}_$(dpkg --print-architecture).deb" \
-    && apt-get install -yq --no-install-recommends ./code-server.deb && rm -f ./code-server.deb \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-
 USER $NB_UID

-RUN mamba install -c nvidia -c conda-forge --yes \
+RUN mamba install -c conda-forge --yes \
        'curl' \
        'ghostscript' \
        'graphviz' \
+        'jq' \
        'rasterio' \
        'gdal' \
        'libgdal-arrow-parquet' \
-    && mamba clean --all -f -y \
-    && pip install --no-cache-dir --extra-index-url="https://pypi.nvidia.com" --extra-index-url "https://download.pytorch.org/whl/cu124" \
+    mamba install -c pytorch -c nvidia --yes \
+        # PyTorch with CUDA, upgrade pytorch-cuda when available
+        'pytorch' \
+        'torchvision' \
+        'torchaudio' \
+        'pytorch-cuda=12.4' \
+    mamba install -c conda-forge -c nvidia --yes \
        # NVIDIA cuQuantum SDK
-        'cuquantum-cu12' \
-        'cuquantum-python-cu12' \
-        # PyTorch with CUDA, upgrade index-url when available
-        torch \
-        torchvision \
-        torchaudio \
-        # Tensorflow with CUDA
-        jupyter-server-proxy \
-        tensorflow \
-        # JAX and Flax with CUDA
-        'jax[cuda12]' \
-        flax \
+        'cuquantum' \
+        'cuquantum-python' \
+        'cuda-version=12' \
+    && mamba clean --all -f -y \
+    && pip install --no-cache-dir --extra-index-url="https://pypi.nvidia.com" \
+        "tensorflow[and-cuda]==$(python3 -c 'import tensorflow as tf; print(tf.__version__)' 2>/dev/null)" \
        # Keras and other packages
        keras \
        keras-hub \
        keras-nlp \
        keras-tuner \
+        # JAX and Flax with CUDA
+        'jax[cuda12]' \
+        flax \
        # https://github.com/explosion/cython-blis/issues/117#issuecomment-2599094860
        'spacy<3.8.0' \
        accelerate \
@@ -139,6 +62,15 @@ RUN mamba install -c nvidia -c conda-forge --yes \
    && fix-permissions "${CONDA_DIR}" \
    && fix-permissions "/home/${NB_USER}"

-COPY --chown="${NB_UID}:${NB_GID}" 20tensorboard-proxy-env.sh /usr/local/bin/before-notebook.d/
+USER root
+
+# Install code-server
+RUN CODESERVER_VERSION="$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | jq -r '.tag_name' | sed 's/[^0-9\.\-]*//g')" \
+    && curl -o code-server.deb -fsSL "https://github.com/coder/code-server/releases/download/v${CODESERVER_VERSION}/code-server_${CODESERVER_VERSION}_$(dpkg --print-architecture).deb" \
+    && apt-get -q update \
+    && apt-get install -yq --no-install-recommends ./code-server.deb && rm -f ./code-server.deb \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+USER $NB_UID

 WORKDIR $HOME