# 使用官方 PyTorch 镜像作为基础镜像
FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel

# 更新包列表并安装系统级软件包
RUN apt-get update && \
    apt-get install -y rclone curl vim iputils-ping sudo && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# 设置环境变量，避免在非交互模式下使用 conda 时出现警告
ENV CONDA_AUTO_UPDATE_CONDA=false \
    PATH=/opt/conda/envs/py310/bin:$PATH

# 创建一个新的 conda 环境并激活它
RUN conda create -n py310 python=3.10 -y && \
    echo "source activate py310" > ~/.bashrc

# 安装所需的 Python 包和 Jupyter Notebook
RUN conda run -n py310 pip install --no-cache-dir \
    transformers==4.41.2 \
    deepspeed==0.14.3 \
    peft==0.11.1 \
    torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121 \
    numpy==1.23.5 \
    aim==3.19.1 \
    jupyter \
    -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn

# 暴露 Jupyter 默认端口
EXPOSE 8888

# 设置工作目录
WORKDIR /workspace


# 启动 Jupyter Notebook 服务
CMD ["conda", "run", "-n", "py310", "jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root"]
