FROM nvidia/cuda:12.1.0-devel-ubuntu22.04

# Evitar prompts interactivos
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# Instalar dependencias del sistema
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    && rm -rf /var/lib/apt/lists/*

# Crear directorio de trabajo
WORKDIR /app

# Actualizar pip
RUN pip3 install --upgrade pip setuptools wheel

# Instalar dependencias de PyTorch (CUDA 12.1)
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Instalar ExLlamaV2
RUN pip3 install exllamav2

# Instalar otras dependencias
RUN pip3 install \
    psycopg2-binary \
    huggingface-hub \
    sentencepiece \
    ninja

# Instalar python-is-python3 para compatibilidad
RUN apt-get update && apt-get install -y python-is-python3 && rm -rf /var/lib/apt/lists/*

# Copiar código del worker
COPY workers/llm_categorizer_worker.py /app/workers/llm_categorizer_worker.py
COPY workers/__init__.py /app/workers/__init__.py

# Crear directorios para modelos y cache
RUN mkdir -p /app/models/llm /app/hf_cache

# Variables de entorno
ENV HF_HOME=/app/hf_cache
ENV TRANSFORMERS_CACHE=/app/hf_cache

# Healthcheck opcional
HEALTHCHECK --interval=60s --timeout=10s --start-period=120s \
    CMD python3 -c "import sys; sys.exit(0)" || exit 1

# Comando por defecto
CMD ["python3", "-m", "workers.llm_categorizer_worker"]