FROM python:3.11-slim

WORKDIR /app

# System deps
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential git && \
    rm -rf /var/lib/apt/lists/*

# Writable dirs for matplotlib and HF cache
ENV MPLCONFIGDIR=/tmp/matplotlib
ENV HF_HOME=/tmp/hf_cache
ENV TRITON_CACHE_DIR=/tmp/triton_cache

# Stage 1: Install torch first (flash-attn needs it at build time)
RUN pip install --no-cache-dir torch==2.4.1

# Stage 2: Install everything else
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# App code
COPY . .

# Create writable dirs AFTER copy
RUN mkdir -p /app/models /app/outputs /app/logs /app/eval_logs /tmp/matplotlib /tmp/hf_cache /tmp/triton_cache && \
    chmod -R 777 /app/models /app/outputs /app/logs /app/eval_logs /tmp/matplotlib /tmp/hf_cache /tmp/triton_cache /app

EXPOSE 7860

CMD ["python", "app.py"]