Dockerfile.pytorch-2.3.1_cuda-12.1_python-3.11_ubuntu-22.04

FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04

ARG PYTHON_VERSION=3.11.9
ARG TORCH_VERSION=2.3.1
ARG TORCHVISION_VERSION=0.18.1
ARG TORCHAUDIO_VERSION=2.3.1
ARG CUDA_VERSION=cu121

ENV LANG=C.UTF-8

# 更新软件源，并安装基本工具和依赖
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    gcc \
    g++ \
    make \
    wget \
    curl \
    openssh-server \
    sudo \
    bzip2 \
    ca-certificates \
    git \
    vim \
    libgl1 libsm6 libxext6 libglib2.0-0 \
    language-pack-en language-pack-zh* \
    screen tmux unzip && \
    apt-get clean && rm -r /var/lib/apt/lists/*

# 设置工作目录为 /root
WORKDIR /root

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

# 安装 Miniconda
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
    bash Miniconda3-latest-Linux-x86_64.sh -b -p /root/miniconda3 && \
    rm Miniconda3-latest-Linux-x86_64.sh

# 配置 conda 环境
ENV PATH=/root/miniconda3/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
RUN source /root/miniconda3/bin/activate && \
    conda init --all && \ 
    source /root/.bashrc
RUN conda remove -y --name base conda-anaconda-telemetry && \
    conda install -y python=${PYTHON_VERSION} && \
    conda clean -y --all

# 安装 PyTorch
RUN pip3 install torch==${TORCH_VERSION}+${CUDA_VERSION} torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION} torchaudio==${TORCHAUDIO_VERSION}+${CUDA_VERSION} --index-url https://download.pytorch.org/whl/${CUDA_VERSION} && \
    rm -r ~/.cache/pip/*

# 安装 Jupyter Notebook
RUN pip install notebook jupyterlab-language-pack-zh-CN && \
    jupyter notebook || true && \
    jupyter notebook --generate-config && \
    rm -r ~/.cache/pip/*


COPY jupyter_notebook_config.py /root/.jupyter/jupyter_notebook_config.py
COPY jupyter /etc/init.d/jupyter
RUN chmod +x /etc/init.d/jupyter

# 配置 SSH 允许 root 登录和密码认证
RUN sed -i 's/^#PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config && \
    sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication yes/' /etc/ssh/sshd_config

# 创建 SSH 运行时目录
RUN mkdir -p /var/run/sshd

# 安装 Prometheus Node-Exporter
RUN apt-get update && apt-get install -y --no-install-recommends prometheus prometheus-node-exporter && \
    apt-get clean && rm -r /var/lib/apt/lists/*
COPY prometheus /etc/init.d/prometheus
COPY prometheus.yml /etc/prometheus/prometheus.yml
RUN chmod +x /etc/init.d/prometheus

COPY prometheus-node-exporter /etc/init.d/prometheus-node-exporter
RUN chmod +x /etc/init.d/prometheus-node-exporter

# 安装 Nvidia-Gpu-Exporter
RUN wget https://github.com/utkuozdemir/nvidia_gpu_exporter/releases/download/v1.2.1/nvidia_gpu_exporter_1.2.1_linux_x86_64.tar.gz -O nvidia_gpu_exporter_1.2.1_linux_x86_64.tar.gz && \
    tar -zxvf nvidia_gpu_exporter_1.2.1_linux_x86_64.tar.gz && \
    mv nvidia_gpu_exporter /usr/bin/ && \
    rm nvidia_gpu_exporter_1.2.1_linux_x86_64.tar.gz LICENSE
COPY nvidia_gpu_exporter /etc/init.d/nvidia_gpu_exporter
RUN chmod +x /etc/init.d/nvidia_gpu_exporter

# 安装 Grafana
RUN wget -q -O - https://packages.grafana.com/gpg.key | apt-key add - && \
    echo "deb https://packages.grafana.com/oss/deb stable main" | tee -a /etc/apt/sources.list.d/grafana.list && \
    apt-get update && apt-get install -y --no-install-recommends grafana && \
    apt-get clean && rm -r /var/lib/apt/lists/*
COPY grafana.db /var/lib/grafana/grafana.db
RUN chown grafana:grafana /var/lib/grafana/grafana.db

RUN apt-get clean || true && \ 
    rm -r /var/lib/apt/lists/* || true && \
    conda clean -y --all || true  && \
    rm -r ~/.cache/pip/* || true 

# 暴露 Jupyter Notebook 的默认端口 8888 和 SSH 默认端口 22
EXPOSE 8888 22 3000

RUN echo -e 'export PATH="/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH"\nexport LANG="C.UTF-8"\n/etc/update-motd.d/10-help-text' >> /root/.bashrc

# 设置 MOTD
COPY 10-help-text /etc/update-motd.d/10-help-text
RUN chmod +x /etc/update-motd.d/10-help-text

# 换源
RUN sed -i 's/archive.ubuntu.com/mirrors.ustc.edu.cn/g' /etc/apt/sources.list
COPY .condarc /root/.condarc
RUN pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple/

# 修改启动脚本
RUN sed -i ':a;$!{N;ba}; s|\necho\n|\necho\nrm -r /var/lib/prometheus/metrics2/*\nulimit -n 64000\nservice ssh start\nservice jupyter start\nservice prometheus start\nservice prometheus-node-exporter start\nservice nvidia_gpu_exporter start\nservice grafana-server start|' /opt/nvidia/nvidia_entrypoint.sh