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 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 


RUN if [ $BASE_IMAGE_IS_TORCH -eq 0 ];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 [ $BASE_IMAGE_IS_TORCH -eq 0 && -n "$TORCH_VERSION" ];then \
    pip3 install torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION torchaudio==$TORCHAUDIO_VERSION \
    --index-url https://download.pytorch.org/whl/cu$(echo "$BASE_IMAGE" | awk -F'[:-]' '{n=split($2,a,"."); print a[1] a[2]}') \
    && 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 ) && \
    pip install --no-cache-dir tensorflow[and-cuda]==$TENSORFLOW_VERSION \
    tensorflow-text==$tf_version_minor.* tf-models-official==$tf_version_minor.* && \
    apt-get update -y && \
    apt-get install --no-install-recommends -y libnvinfer8 libnvjitlink-12-3 libnvjpeg-12-3 libnvinfer-plugin8; 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_dataset-0.1.0-py3-none-any.whl /tmp/jupyter-extension/
COPY ./jupyter_ext_logo-0.1.0-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 ./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