# If the device is Battlemage, we need to set UBUNTU_VERSION to 24.10

# Usage: docker build --build-arg UBUNTU_VERSION=24.04 --build-arg PYTHON_VERSION=3.10 -t sglang:xpu_kernel -f  Dockerfile.xpu --no-cache .

# Use Intel deep learning essentials base image with Ubuntu 24.04
FROM intel/deep-learning-essentials:2025.1.3-0-devel-ubuntu24.04

# Avoid interactive prompts during package install
ENV DEBIAN_FRONTEND=noninteractive

# Define build arguments
ARG PYTHON_VERSION=3.10

ARG SG_LANG_REPO=https://github.com/sgl-project/sglang.git
ARG SG_LANG_BRANCH=main

ARG SG_LANG_KERNEL_REPO=https://github.com/sgl-project/sgl-kernel-xpu.git
ARG SG_LANG_KERNEL_BRANCH=main

RUN useradd -m -d /home/sdp -s /bin/bash sdp && \
    chown -R sdp:sdp /home/sdp

# Switch to non-root user 'sdp'
USER sdp

# Set HOME and WORKDIR to user's home directory
ENV HOME=/home/sdp
WORKDIR /home/sdp

RUN curl -fsSL -v -o miniforge.sh -O https://github.com/conda-forge/miniforge/releases/download/25.1.1-0/Miniforge3-Linux-x86_64.sh && \
    bash miniforge.sh -b -p ./miniforge3 && \
    rm miniforge.sh && \
    # Initialize conda environment and install pip
    . ./miniforge3/bin/activate && \
    conda create -y -n py${PYTHON_VERSION} python=${PYTHON_VERSION} && \
    conda activate py${PYTHON_VERSION} && \
    conda install pip && \
    # Append environment activation to .bashrc for interactive shells
    echo ". /home/sdp/miniforge3/bin/activate; conda activate py${PYTHON_VERSION}; . /opt/intel/oneapi/setvars.sh; cd /home/sdp" >> /home/sdp/.bashrc

USER root
RUN apt-get update && apt install -y intel-ocloc

# Switch back to user sdp
USER sdp

RUN --mount=type=secret,id=github_token \
    cd /home/sdp && \
    . /home/sdp/miniforge3/bin/activate && \
    conda activate py${PYTHON_VERSION} && \
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu

RUN --mount=type=secret,id=github_token \
    cd /home/sdp && \
    . /home/sdp/miniforge3/bin/activate && \
    conda activate py${PYTHON_VERSION} && \
    echo "Cloning ${SG_LANG_BRANCH} from ${SG_LANG_REPO}" && \
    git clone --branch ${SG_LANG_BRANCH} --single-branch ${SG_LANG_REPO} && \
    cd sglang && cd python && \
    cp pyproject_xpu.toml pyproject.toml && \
    pip install . && \
    echo "Cloning ${SG_LANG_KERNEL_REPO} from ${SG_LANG_KERNEL_BRANCH}" && \
    git clone --branch ${SG_LANG_KERNEL_BRANCH} --single-branch ${SG_LANG_KERNEL_REPO} && \
    cd sgl-kernel-xpu && \
    pip install -v . && \
    pip install msgspec blake3 py-cpuinfo compressed_tensors gguf partial_json_parser einops --root-user-action=ignore && \
    pip uninstall pytorch-triton-xpu -y && \
    pip install --pre pytorch-triton-xpu --index-url https://download.pytorch.org/whl/xpu && \
    conda install libsqlite=3.48.0 -y && \
    # Add environment setup commands to .bashrc again (in case it was overwritten)
    echo ". /home/sdp/miniforge3/bin/activate; conda activate py${PYTHON_VERSION}; cd /home/sdp" >> /home/sdp/.bashrc

# Use bash as default shell with initialization from .bashrc
SHELL ["bash", "-c"]

# Start an interactive bash shell with all environment set up
USER sdp
CMD ["bash", "-c", "source /home/sdp/.bashrc && exec bash"]