Dockerfile.jupyterlab_ubuntu

ARG BASE_IMAGE="nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04"

FROM $BASE_IMAGE
ARG BASE_IMAGE
ARG DEBIAN_FRONTEND=noninteractive
LABEL module="jupyter"

# ----- torch args -----
# 是否基于torch镜像构建
ARG BASE_IMAGE_IS_TORCH=0
ARG TORCH_VERSION
ARG TORCHVISION_VERSION
ARG TORCHAUDIO_VERSION

# ----- tensorflow args -----
ARG TENSORFLOW_VERSION
ARG IMAGE_TAG

#ARG CONDA_URL="https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-py310_24.7.1-0-Linux-x86_64.sh"
# ----- paddlepaddle args -----
ARG PADDLEPADDLE_VERSION
ARG PADDLENLP_VERSION
ARG CUDA_VERSION
ARG PADDLE_URL
# ----- conda and python args ----
ARG CONDA_URL
#ARG CONDA_URL="https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-py310_24.7.1-0-Linux-x86_64.sh"
ARG SOURCES="-i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn"
ENV TZ=Asia/Shanghai
ENV SHELL=/bin/bash \
    LC_ALL=en_US.UTF-8 \
    LANG=en_US.UTF-8 \
    LANGUAGE=en_US.UTF-8
ENV HF_ENDPOINT=https://hf-mirror.com
SHELL ["/bin/bash","-c"]
USER root

RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak \
    && mv /etc/apt/sources.list.d /etc/apt/sources.list.d.bak \
    && version_id=$(cat /etc/os-release | grep -i 'VERSION_CODENAME' | awk -F= '{print $2}') \
    && echo -e "deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id} main restricted universe multiverse \n\
    deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-security main restricted universe multiverse\n\
    deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-updates main restricted universe multiverse\n\
    deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-proposed main restricted universe multiverse\n\
    deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-backports main restricted universe multiverse\n\
    deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id} main restricted universe multiverse\n\
    deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-security main restricted universe multiverse\n\
    deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-updates main restricted universe multiverse\n\
    deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-proposed main restricted universe multiverse\n\
    deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ ${version_id}-backports main restricted universe multiverse" > /etc/apt/sources.list \
    && apt-get update -y \
    && apt-get install --no-install-recommends -y vim openssl libssl-dev net-tools locales wget git git-lfs sudo openssh-client openssh-server \
    && locale-gen en_US.UTF-8 

# ----- conda and python install -----

RUN if [ -n "$CONDA_URL" ];then \
    mkdir -p /tmp/conda-extension \
    && cd /tmp/conda-extension \
    && wget $CONDA_URL \
    && bash $(echo $CONDA_URL | awk -F "/" '{print $NF}') -b -p /opt/conda \
    && echo "export PATH=\$PATH:/opt/conda/bin" >> /etc/profile.d/sothisai.sh \
    && cd .. \
    && rm -rf /tmp/conda-extension; fi

ENV PATH=$PATH:/opt/conda/bin

RUN pip3 install --upgrade pip ${SOURCES} || pip install --upgrade pip ${SOURCES} \
    && pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
    && mv /etc/apt/sources.list.bak /etc/apt/sources.list \
    && mv /etc/apt/sources.list.d.bak /etc/apt/sources.list.d

RUN if [ -n "$TORCH_VERSION" ] && [ $BASE_IMAGE_IS_TORCH -eq 0 ];then \
    cuda_version=$(echo "$BASE_IMAGE" | awk -F'[:-]' '{n=split($2,a,"."); print a[1] a[2]}') \
    && if [ "${TORCH_VERSION%%.*}" -eq 1 ]; then \
    pip3 install --no-cache-dir "numpy<2" \
    && pip3 install --no-cache-dir  torch==$TORCH_VERSION+cu$cuda_version torchvision==$TORCHVISION_VERSION+cu$cuda_version torchaudio==$TORCHAUDIO_VERSION \
    --extra-index-url https://download.pytorch.org/whl/cu$cuda_version;else \
    pip3 install --no-cache-dir torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION torchaudio==$TORCHAUDIO_VERSION \
    --index-url https://download.pytorch.org/whl/cu$cuda_version; fi \
    && rm -r /root/.cache/pip; fi

RUN if [ -n "$TORCH_VERSION" ];then \
    pip install --no-cache-dir transformers accelerate diffusers; fi

RUN if [ -n "$TENSORFLOW_VERSION" ]; then \
    tf_version_minor=$(echo $TENSORFLOW_VERSION | cut -d'.' -f1-2 ); \
    [ "$tf_version_minor" == "2.13" ] || [ "$tf_version_minor" == "2.18" ] && tensorflow_text_version=$tf_version_minor.0rc0 || tensorflow_text_version=$tf_version_minor.*; \
    pip install --no-cache-dir tensorflow[and-cuda]==$TENSORFLOW_VERSION tensorflow-text==$tensorflow_text_version tensorflow-hub; \
    # 2.16.1必须手动添加环境变量
    if [ $TENSORFLOW_VERSION == 2.16.1 ]; then \
    python_version=$(echo $IMAGE_TAG | awk -F'[-:]' '{for(i=3;i<=NF;i++) if($i ~ /^py[0-9]+\.[0-9]+$/) {gsub(/^py/,"",$i); print $i; exit}}') && \
    CUDNN_PATH=/opt/conda/lib/python$python_version/site-packages/nvidia/cudnn && \
    echo "export CUDNN_PATH=$CUDNN_PATH" >> /etc/bash.bashrc && \
    echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDNN_PATH/lib:/usr/local/cuda/lib64" >> /etc/bash.bashrc; \
    fi; fi

# ----- paddlepaddle install -----
RUN if [ -n "$PADDLEPADDLE_VERSION" ] && [ -n "$PADDLE_URL" ]; then \
    pip install paddlepaddle-gpu==$PADDLEPADDLE_VERSION -f $PADDLE_URL -i $PADDLE_URL \
    && rm -r /root/.cache/pip; \
fi

RUN if [ -n "$PADDLEPADDLE_VERSION" ] && [ -z "$PADDLE_URL" ]; then \
    pip install paddlepaddle-gpu==$PADDLEPADDLE_VERSION -i https://pypi.tuna.tsinghua.edu.cn/simple \
    && rm -r /root/.cache/pip; \
fi

RUN if [ -n "$PADDLENLP_VERSION" ] ; then \
    pip install paddlenlp==$PADDLENLP_VERSION ppdiffusers huggingface_hub --no-cache-dir  -i https://pypi.tuna.tsinghua.edu.cn/simple && \
    pip install --upgrade ppdiffusers --no-deps  && rm -r /root/.cache/pip; \
fi


COPY ./python-requirements.txt /tmp/
RUN pip install --no-cache-dir -r /tmp/python-requirements.txt

### 添加HF的镜像变量
ENV HF_ENDPOINT=https://hf-mirror.com
RUN echo 'export HF_ENDPOINT=https://hf-mirror.com' >> /etc/bash.bashrc && \
    echo 'export HF_ENDPOINT=https://hf-mirror.com' >> ~/.bashrc && \
    echo 'export CUDA_HOME=/usr/local/cuda' >> /etc/bash.bashrc

# 安装jupyterlab插件, 生成sshd配置
RUN pip3 install jupyterlab==4.0.0
RUN mkdir -p /tmp/jupyter-extension
COPY ./jupyter_ext_model-1.0.7-py3-none-any.whl /tmp/jupyter-extension/
COPY ./jupyter_ext_platform-1.0.11-py3-none-any.whl /tmp/jupyter-extension/
COPY ./jupyterlab_language_pack_zh_cn-4.0.post6-py2.py3-none-any.whl /tmp/jupyter-extension/
COPY ./git_lfs-1.6-py2.py3-none-any.whl /tmp/jupyter-extension/
COPY ./requirements.txt /tmp/jupyter-extension/
COPY ./static /tmp/jupyter-extension/static
COPY extension.sh /tmp/
RUN chmod +x /tmp/extension.sh && \
    /tmp/extension.sh

COPY ./jupyter_lab_config.py /etc/
RUN jupytersite="$(python3 -m pip show jupyterlab | grep -i '^location' | awk '{print $2}')" \
    && JUPYTERLAB_DIR=$(echo ${jupytersite%%/lib/python*}) \
    && if [[ ${JUPYTERLAB_DIR}/bin/jupyter != "/opt/conda/bin/jupyter" ]]; \
    then mkdir -p /opt/conda/bin; \
    ln -sf ${JUPYTERLAB_DIR}/bin/jupyter /opt/conda/bin/jupyter;fi \
    && echo "export JUPYTER_CONFIG_PATH=/etc" >> /etc/bash.bashrc \
    && echo "export JUPYTER=${JUPYTERLAB_DIR}/bin/jupyter" >> /etc/bash.bashrc\
    && echo "export PATH=${JUPYTERLAB_DIR}/bin:${PATH}" >> /etc/bash.bashrc \
    && mkdir -p /var/run/sshd \
    && ssh-keygen  -A \
    && sed -i "s/#UseDNS .*/UseDNS no/" /etc/ssh/sshd_config

EXPOSE 8888