FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    TZ=Europe/Paris \
    CONDA_AUTO_UPDATE_CONDA=false \
    HOME=/home/user \
    PATH=/home/user/miniconda/bin:$PATH \
    HF_HOME=/data/huggingface \
    HF_HUB_DISABLE_TELEMETRY=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    GRADIO_ALLOW_FLAGGING=never \
    GRADIO_NUM_PORTS=1 \
    GRADIO_SERVER_NAME=0.0.0.0 \
    GRADIO_THEME=huggingface \
    SYSTEM=spaces \
    SHELL=/bin/bash

# Base utilities
RUN rm -f /etc/apt/sources.list.d/*.list && \
    apt-get update && apt-get install -y --no-install-recommends \
    curl \
    ca-certificates \
    git \
    wget \
    procps \
    git-lfs \
    zip \
    unzip \
    htop \
    vim \
    nano \
    bzip2 \
    libx11-6 \
    build-essential \
    libsndfile-dev \
    software-properties-common \
 && rm -rf /var/lib/apt/lists/*

# nvtop
RUN add-apt-repository ppa:flexiondotorg/nvtop && \
    apt-get update && \
    apt-get install -y --no-install-recommends nvtop && \
    rm -rf /var/lib/apt/lists/*

# Node.js 21
RUN curl -fsSL https://deb.nodesource.com/setup_21.x | bash - && \
    apt-get update && \
    apt-get install -y nodejs && \
    npm install -g configurable-http-proxy && \
    rm -rf /var/lib/apt/lists/*

# Working directory
WORKDIR /app

# User setup
RUN adduser --disabled-password --gecos '' --shell /bin/bash user && \
    mkdir -p /home/user/.cache /home/user/.config /home/user/app && \
    chown -R user:user /home/user

# Miniconda Python 3.12 (latest stable)
USER user
RUN curl -fsSL -o /home/user/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-py312_24.7.1-0-Linux-x86_64.sh && \
    bash /home/user/miniconda.sh -b -p /home/user/miniconda && \
    rm -f /home/user/miniconda.sh && \
    conda install -y python=3.12.13 pip=26.0.1 && \
    conda clean -ya

WORKDIR /home/user/app

# Back to root for system packages / startup
USER root

RUN --mount=target=/root/packages.txt,source=packages.txt \
    apt-get update && \
    xargs -r -a /root/packages.txt apt-get install -y --no-install-recommends && \
    rm -rf /var/lib/apt/lists/*

RUN --mount=target=/root/on_startup.sh,source=on_startup.sh,readwrite \
    bash /root/on_startup.sh

RUN mkdir -p /data/huggingface /workspace && chown -R root:root /data /workspace

# Python packages
ARG FLASH_ATTN_WHEEL_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"
ARG FLASH_ATTN_WHEEL_SHA256="4e2f9e39313266b1544b68138b15b91ee6221eccf14f7902b7c6620351340810"

USER user
RUN pip install --no-cache-dir \
    packaging==26.2 \
    psutil==7.2.2 \
    ninja==1.13.0 \
    wheel==0.46.3 && \
    pip install --no-cache-dir \
    torch==2.9.1 \
    --index-url https://download.pytorch.org/whl/cu128

RUN --mount=target=requirements.txt,source=requirements.txt \
    pip install --no-cache-dir --upgrade -r requirements.txt

RUN FLASH_ATTN_WHEEL="/tmp/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" && \
    curl -fL "$FLASH_ATTN_WHEEL_URL" -o "$FLASH_ATTN_WHEEL" && \
    echo "$FLASH_ATTN_WHEEL_SHA256  $FLASH_ATTN_WHEEL" | sha256sum -c - && \
    pip install --no-cache-dir "$FLASH_ATTN_WHEEL" && \
    rm -f "$FLASH_ATTN_WHEEL"

RUN python -c "import torch, flash_attn; print(f'torch={torch.__version__} cuda={torch.version.cuda} flash_attn={flash_attn.__version__}')"

# App files
USER root
COPY --chown=root:root . /home/user/app

RUN chmod +x /home/user/app/start_server.sh

# Jupyter template path for Python 3.12
COPY --chown=root:root login.html /home/user/miniconda/lib/python3.12/site-packages/jupyter_server/templates/login.html

CMD ["./start_server.sh"]