FROM huggingface/transformers-pytorch-gpu:latest

WORKDIR /app

# Instalar dependências do sistema
# python-is-python3 cria automaticamente o symlink python -> python3
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    python3 \
    python3-pip \
    python-is-python3 \
    && rm -rf /var/lib/apt/lists/*

# Verificar que python está disponível (entrypoint do NVIDIA precisa)
RUN python --version && \
    python3 --version && \
    echo "✅ Python disponível: $(which python)"

# Instalar dependências Python (constraint em torch = versão da imagem base, evita upgrade que quebra torchvision)
COPY requirements.txt .
RUN python3 -c "import torch; open('/tmp/constraints.txt','w').write('torch==' + torch.__version__.split('+')[0] + '\n')" && \
    python3 -m pip install --no-cache-dir --upgrade pip && \
    python3 -m pip install --no-cache-dir -r requirements.txt -c /tmp/constraints.txt && \
    python3 -m pip install --no-cache-dir --upgrade torchvision && \
    (python3 -m pip install --no-cache-dir nvidia-cuda-nvjitlink-cu12 || true) && \
    python3 -m pip uninstall -y bitsandbytes 2>/dev/null || true

# Treino LoRA fp16 sem quantizacao: sem bitsandbytes (imagem base pode traze-lo; libnvJitLink incompativel)
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/nvidia/cuda_nvjitlink/lib:/usr/local/lib/python3.10/dist-packages/nvidia/cuda_nvjitlink/lib/x86_64-linux-gnu:/usr/local/lib/python3.11/dist-packages/nvidia/cuda_nvjitlink/lib:/usr/local/lib/python3.11/dist-packages/nvidia/cuda_nvjitlink/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}

# Copiar scripts de treinamento
COPY training_env.py /app/training_env.py
COPY reporting.py /app/reporting.py
COPY space_lock.py /app/space_lock.py
COPY train.py /app/train.py
COPY app.py /app/app.py
COPY benchmarks/gold_sample.json /app/benchmarks/gold_sample.json

# Criar diretório de logs
RUN mkdir -p /app/logs

# Configurar variáveis de ambiente padrão (podem ser sobrescritas)
ENV MODEL_NAME=Qwen/Qwen2.5-1.5B-Instruct
ENV DATASET_REPO=beAnalytic/eda-training-dataset
ENV OUTPUT_REPO=beAnalytic/eda-llm-qwen2.5-lora
ENV OMP_NUM_THREADS=1
ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# T4 ~15GB: seq 256 por defeito (OOM com varios processos ou logits na loss); sobrescreve no Space
ENV MAX_SEQ_LENGTH=256
ENV PER_DEVICE_TRAIN_BATCH_SIZE=1
ENV PER_DEVICE_EVAL_BATCH_SIZE=1
ENV GRADIENT_ACCUMULATION_STEPS=8

# Executar treinamento
# Usar 'python' (que será o symlink para python3 criado acima)
# O entrypoint do NVIDIA espera 'python' estar disponível
CMD ["python", "/app/app.py"]