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

FROM $BASE_IMAGE
ARG BASE_IMAGE
LABEL module="vscode"
ARG DEBIAN_FRONTEND=noninteractive
# 是否基于torch镜像构建
ARG BASE_IMAGE_IS_TORCH=0
ARG TORCH_VERSION="2.0.1"
ARG TORCHVISION_VERSION="0.15.2"
ARG TORCHAUDIO_VERSION="2.0.2"
ARG CONDA_URL="https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-py310_24.7.1-0-Linux-x86_64.sh"

ENV TZ=Asia/Shanghai
ENV SHELL=/bin/bash \
    LC_ALL=en_US.UTF-8 \
    LANG=en_US.UTF-8 \
    LANGUAGE=en_US.UTF-8
ARG SOURCES="-i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn"
SHELL ["/bin/bash","-c"]
USER root
WORKDIR /tmp

ENV HF_ENDPOINT=https://hf-mirror.com

COPY ./python-requirements.txt /tmp/

COPY ./code-server_4.22.1_amd64.deb ./
RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak \
&& mv /etc/apt/sources.list.d /etc/apt/sources.list.d.bak \
&& version_codename=$(cat /etc/os-release | grep -i 'VERSION_CODENAME' | awk -F= '{print $2}') \
&& echo -e "deb http://mirrors.aliyun.com/ubuntu/ ${version_codename} main restricted universe multiverse \n\
deb http://mirrors.aliyun.com/ubuntu/ ${version_codename}-security main restricted universe multiverse\n\
deb http://mirrors.aliyun.com/ubuntu/ ${version_codename}-updates main restricted universe multiverse\n\
deb http://mirrors.aliyun.com/ubuntu/ ${version_codename}-proposed main restricted universe multiverse\n\
deb http://mirrors.aliyun.com/ubuntu/ ${version_codename}-backports main restricted universe multiverse\n\
deb-src http://mirrors.aliyun.com/ubuntu/ ${version_codename} main restricted universe multiverse\n\
deb-src http://mirrors.aliyun.com/ubuntu/ ${version_codename}-security main restricted universe multiverse\n\
deb-src http://mirrors.aliyun.com/ubuntu/ ${version_codename}-updates main restricted universe multiverse\n\
deb-src http://mirrors.aliyun.com/ubuntu/ ${version_codename}-proposed main restricted universe multiverse\n\
deb-src http://mirrors.aliyun.com/ubuntu/ ${version_codename}-backports main restricted universe multiverse" > /etc/apt/sources.list \
    && apt-get update -y \
    && apt-get install --no-install-recommends -y \
    openssl \
    libssl-dev \
    net-tools \
    locales \
    wget \
    git \
    sudo \
    openssh-client \
    openssh-server \
    && locale-gen en_US.UTF-8 \
    && dpkg -i code-server_4.22.1_amd64.deb \
    #&& apt-get clean \
    #&& rm -rf /var/lib/apt/lists/* \
    #&& rm -rf /tmprm \
    #&& mv /etc/apt/sources.list.bak /etc/apt/sources.list \
    #&& mv /etc/apt/sources.list.d.bak /etc/apt/sources.list.d \
    && mkdir -p /var/run/sshd \
    && ssh-keygen  -A \
    && sed -i "s/#UseDNS .*/UseDNS no/" /etc/ssh/sshd_config

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

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


EXPOSE 8080