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

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PYTHONPATH=/home/user/app
ENV HF_HOME=/tmp/.cache/huggingface
ENV PIP_NO_CACHE_DIR=1

RUN apt-get update -y && \
    apt-get install -y python3 python3-pip python3-venv git curl && \
    python3 -m pip install --upgrade pip && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

RUN useradd -m -u 1000 user
USER user
ENV HOME=/home/user
ENV PATH=/home/user/.local/bin:$PATH
WORKDIR /home/user/app

# Install torch first (heaviest, cached separately)
RUN pip install torch==2.5.1+cu121 --index-url https://download.pytorch.org/whl/cu121

# Install unsloth's official Colab-compatible dependency bundle.
# This is the ONLY combination unsloth officially supports and tests.
RUN pip install "unsloth[colab-new]"

# Install unsloth core (no-deps to not override colab-new pins)
RUN pip install --no-deps "unsloth @ git+https://github.com/unslothai/unsloth.git"

# Install our additional deps (server + OpenEnv + matplotlib)
RUN pip install \
    flask \
    flask-cors \
    fastapi \
    uvicorn \
    pydantic \
    requests \
    openenv-core \
    PyYAML \
    matplotlib

# Verify non-GPU imports work
RUN python3 -c "import torch; print(f'torch={torch.__version__}')" && \
    python3 -c "import transformers; print(f'transformers={transformers.__version__}')" && \
    python3 -c "import trl; print(f'trl={trl.__version__}')" && \
    python3 -c "import datasets; print(f'datasets={datasets.__version__}')"

COPY --chown=user . /home/user/app

RUN pip install --no-deps -e /home/user/app

RUN python3 -m training.generate_warmup_traces

EXPOSE 7860

# The HF Space receives entrypoint.sh at repo root (promoted by tools/upload_all.py),
# but if someone builds locally from `deploy/training/` it's one directory up.
RUN if [ -f /home/user/app/entrypoint.sh ]; then \
        chmod +x /home/user/app/entrypoint.sh; \
    elif [ -f /home/user/app/deploy/training/entrypoint.sh ]; then \
        cp /home/user/app/deploy/training/entrypoint.sh /home/user/app/entrypoint.sh && \
        chmod +x /home/user/app/entrypoint.sh; \
    fi

CMD ["/home/user/app/entrypoint.sh"]