Dockerfile-notebook

FROM truefoundrycloud/jupyter:0.2.19-sudo
ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
ENV DEBIAN_FRONTEND=noninteractive
USER root
RUN add-apt-repository ppa:flexiondotorg/nvtop -y && \
    apt update && \
    apt install -y --no-install-recommends git curl wget htop nvtop && \
    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb -O /tmp/cuda-keyring_1.1-1_all.deb && \
    dpkg -i /tmp/cuda-keyring_1.1-1_all.deb && \
    apt update && \
    apt install -y --no-install-recommends cuda-toolkit-12-1 libcudnn8=8.9.7.29-1+cuda12.2 libcudnn8-dev=8.9.7.29-1+cuda12.2 libnccl2 libnccl-dev
USER jovyan
COPY requirements.txt notebook-requirements.txt /tmp/llm-finetune/
RUN pip install -U pip wheel setuptools && \
    pip uninstall -y mlflow axolotl && \
    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
USER root
RUN mkdir -p /packages && \
    chown -R jovyan:users /packages
USER jovyan
RUN cd /packages && \
    git clone https://github.com/truefoundry/axolotl && \
    cd axolotl/ && \
    git checkout bfe2d6b2e5019a8d4486ac8a7726089928e6b3a6
RUN cd /packages/axolotl/ && \
    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation -e .[flash-attn,mamba-ssm,fused-dense-lib] && \
    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
COPY --chown=jovyan:users . /tmp_home/jovyan/llm-finetune/
ENV JUPYTER_APP_LAUNCHER_PATH=/home/jovyan/llm-finetune/.jp_app_launcher/