FROM 353750902984.dkr.ecr.us-east-1.amazonaws.com/thirdparty-chainguard-python310:latest-fips-dev

ARG DEBIAN_FRONTEND=noninteractive
ARG CUDA_MAJOR_VERSION=13
ARG CUDA_MINOR_VERSION=0

ENV NVIDIA_DRIVER_CAPABILITIES="compute,utility"
ENV NVIDIA_VISIBLE_DEVICES="all"

USER root

RUN apk update \
    && apk add wget \
    && wget -O /etc/apk/keys/chainguard-extras.rsa.pub https://packages.cgr.dev/extras/chainguard-extras.rsa.pub \
    && echo "https://packages.cgr.dev/extras" | tee -a /etc/apk/repositories \
    && apk update \
    && apk add --no-cache \
    nvidia-cuda-cudart-${CUDA_MAJOR_VERSION}.${CUDA_MINOR_VERSION} \
    nvidia-cuda-nvcc-${CUDA_MAJOR_VERSION}.${CUDA_MINOR_VERSION} \
    make \
    curl \
    git \
    patch

WORKDIR /workspace

ENV CUDA_HOME=/usr/local/cuda-${CUDA_MAJOR_VERSION}.${CUDA_MINOR_VERSION}
ENV PATH=$CUDA_HOME/bin:$PATH
ENV LD_LIBRARY_PATH=$CUDA_HOME/lib64

# Install uv and python dependencies
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
RUN --mount=type=bind,src=pyproject.toml,target=pyproject.toml \
    --mount=type=bind,src=uv.lock,target=uv.lock \
    /root/.local/bin/uv sync --frozen --no-cache --no-dev

# Apply source patches to installed third-party packages. See patches/ for why
# each one exists; they are removed once the fix is released upstream.
RUN --mount=type=bind,src=patches,target=/tmp/patches \
    SITE_PACKAGES=$(/workspace/.venv/bin/python -c "import sysconfig; print(sysconfig.get_paths()['purelib'])") \
    && for p in /tmp/patches/*.patch; do \
        [ -e "$p" ] || continue; \
        echo "Applying $(basename "$p") ..."; \
        patch -d "$SITE_PACKAGES" -p1 --forward -r - < "$p" \
            || { [ $? -eq 1 ] && echo "  (already applied, skipping)"; }; \
    done

# Add the venv to the PATH
ENV PATH=/workspace/.venv/bin:$PATH

# We need to create a mount point for the user to mount their volume
# All persistent data lives in /mount
RUN mkdir -p /mount && chown -R nonroot:nonroot /mount
ENV H2O_LLM_STUDIO_WORKDIR=/mount

# Download the demo datasets and place in the /workspace/demo directory
# Set the environment variable for the demo datasets
ENV H2O_LLM_STUDIO_DEMO_DATASETS=/workspace/demo
COPY --chown=nonroot:nonroot ./llm_studio/download_default_datasets.py /workspace/
RUN python download_default_datasets.py

COPY --chown=nonroot:nonroot ./llm_studio /workspace/llm_studio
COPY --chown=nonroot:nonroot ./prompts /workspace/prompts
COPY --chown=nonroot:nonroot ./model_cards /workspace/model_cards
COPY --chown=nonroot:nonroot ./LICENSE /workspace/LICENSE
COPY --chown=nonroot:nonroot ./entrypoint.sh /workspace/entrypoint.sh
COPY --chown=nonroot:nonroot ./pyproject.toml /workspace/pyproject.toml

ENV HF_HOME=/mount/huggingface
ENV TRITON_CACHE_DIR=/mount/.triton/cache
ENV H2O_WAVE_DATA_DIR=/mount/wave_data
ENV HF_HUB_DISABLE_TELEMETRY=1
ENV DO_NOT_TRACK=1

# Set the environment variables for the wave server
ENV H2O_WAVE_APP_ADDRESS=http://127.0.0.1:8756
ENV H2O_WAVE_MAX_REQUEST_SIZE=25MB
ENV H2O_WAVE_NO_LOG=true
ENV H2O_WAVE_PRIVATE_DIR="/download/@/mount/output/download"

# Make the entrypoint.sh script executable
RUN chmod 755 /workspace/entrypoint.sh

EXPOSE 10101

USER nonroot

ENTRYPOINT [ "/workspace/entrypoint.sh" ]