Dockerfile.gpu

# This Dockerfile is based on this post in the uv GitHub:
# https://github.com/astral-sh/uv/issues/7758#issuecomment-3263282018
FROM python:3.12-slim AS build

# Install 'uv', which we will use to install Python dependencies
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
WORKDIR /app

# uv configuration
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#compiling-bytecode
ENV UV_COMPILE_BYTECODE=1 \
    # Ref: https://docs.astral.sh/uv/guides/integration/docker/#caching
    UV_LINK_MODE=copy \
    # Ref: https://docs.astral.sh/uv/guides/integration/docker/#managing-python-interpreters
    UV_PYTHON_INSTALL_DIR=/opt/python

RUN uv python install 3.12
RUN uv venv
RUN uv pip install --no-cache-dir --upgrade pip setuptools wheel

# Install CPU requirements
COPY requirements.cpu.txt ./
RUN uv pip install --no-cache-dir -r ./requirements.cpu.txt

# Install GPU PyTorch requirements
COPY requirements.torch.gpu.txt ./
RUN uv pip install --no-cache-dir -r ./requirements.torch.gpu.txt

# Download models during build instead of copying from local
COPY scripts/model_download.bash /tmp/model_download.bash
RUN . .venv/bin/activate && \
    uv pip install --no-cache-dir huggingface-hub && \
    bash /tmp/model_download.bash && \
    rm /tmp/model_download.bash

COPY app ./app
COPY main.py ./

# ----------------------------------------------------------------------------

FROM nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04 AS run
WORKDIR /app

ENV PYTHONDONTWRITEBYTECODE="1" \
    PYTHONUNBUFFERED="1" \
    DEBIAN_FRONTEND="noninteractive"

# The container will run as a non-root user with unknown UID and GID, so we
# need to give everyone permission to all of the files
COPY --from=build --chmod=777 /opt/python /opt/python
COPY --from=build --chmod=777 /app /app

# Add python env to PATH and app folder to PYTHONPATH
ENV PATH="/app/.venv/bin:${PATH}" \
    PYTHONPATH="/app"

EXPOSE 8000

# torch.compile tries to cache things in /tmp/torchinductor_{username},
# which doesn't work because we don't have a user account
ENV TORCHINDUCTOR_CACHE_DIR="/tmp/torchinductor"

ENTRYPOINT ["python", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]