# Dockerfile.gpu - DocTR GPU Dockerfile for amd64/arm64
#
# Build:
#   docker build -t doctr-gpu:latest -f Dockerfile.gpu .
#
# Run:
#   docker run --gpus all -p 8003:8000 -v ./dataset:/app/dataset doctr-gpu:latest

# CUDA 13.0 for Blackwell (sm_121) and GH200/GB200 support
FROM nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04

LABEL maintainer="Sergio Jimenez"
LABEL description="DocTR Tuning REST API - GPU/CUDA version"

WORKDIR /app

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV CUDA_VISIBLE_DEVICES=0
ENV DOCTR_DET_ARCH=db_resnet50
ENV DOCTR_RECO_ARCH=crnn_vgg16_bn

# Install Python 3.12 and system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.12 \
    python3.12-venv \
    python3-pip \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender1 \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.12 /usr/bin/python

# Copy requirements first for better caching
COPY requirements.txt .

# Install PyTorch with CUDA support first (cu128 index has amd64 + arm64 wheels)
RUN python -m pip install --no-cache-dir --break-system-packages \
    torch torchvision --index-url https://download.pytorch.org/whl/cu128

# Install remaining dependencies from requirements.txt (skip torch, already installed)
RUN grep -v "^torch" requirements.txt | python -m pip install --no-cache-dir --break-system-packages -r /dev/stdin

# Copy application code
COPY doctr_tuning_rest.py .
COPY dataset_manager.py .

# Volume for dataset and model cache
VOLUME ["/app/dataset", "/root/.cache/doctr"]

# Expose API port
EXPOSE 8000

# Health check (longer start period for model download)
HEALTHCHECK --interval=30s --timeout=10s --start-period=180s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1

# Run the API server
CMD ["uvicorn", "doctr_tuning_rest:app", "--host", "0.0.0.0", "--port", "8000"]
