FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel

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

# Python deps
COPY requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir -r /tmp/requirements.txt packaging

# Copy scripts
WORKDIR /app
COPY prepare_data.py /app/
COPY train_lora.py /app/
COPY run.sh /app/
RUN chmod +x /app/run.sh

# Data and output dirs
RUN mkdir -p /data/processed /data/lora-output /data/models

# Default: run the full pipeline
ENTRYPOINT ["/app/run.sh"]