Dockerfile-notebook

FROM tfy.jfrog.io/tfy-images/jupyter:0.3.8-cu121-py3.11.10-sudo
SHELL ["/bin/bash", "-c"]
USER root
RUN [ "$(/usr/local/cuda/bin/nvcc --version | egrep -o "V[0-9]+\.[0-9]+" | cut -c2-)" = "12.1" ] || (echo "Error: CUDA version is not 12.1" && exit 1)
ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
ENV DEBIAN_FRONTEND=noninteractive
# upgrade libnccl
RUN apt update && \
    apt install -y --no-install-recommends git curl wget && \
    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/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 --allow-change-held-packages libnccl2 libnccl-dev

# Install torch and axolotl
USER jovyan
COPY base-requirements.txt requirements.txt notebook-requirements.txt /tmp/llm-finetune/
RUN pip install -U pip wheel setuptools && \
    pip uninstall -y axolotl torch && \
    pip install -U --no-cache-dir --use-pep517 -r /tmp/llm-finetune/base-requirements.txt && \
    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install --no-cache-dir --no-build-isolation --use-pep517 -r /tmp/llm-finetune/requirements.txt && \
    pip install --no-cache-dir --use-pep517 -r /tmp/llm-finetune/notebook-requirements.txt

# Setup editable packages
USER root
RUN mkdir -p /packages && \
    chown -R jovyan:users /packages

# Install axolotl_truefoundry plugin with our requirements overrides over axolotl
COPY --chown=jovyan:users plugins/axolotl_truefoundry /packages/axolotl_truefoundry
RUN cd /packages/axolotl_truefoundry/ && \
    pip install --no-cache-dir -e .

# Add source code for finetuning
COPY --chown=jovyan:users . /tmp_home/jovyan/llm-finetune/
ENV JUPYTER_APP_LAUNCHER_PATH=/home/jovyan/llm-finetune/.jp_app_launcher/