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

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# Install system essentials
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    python3-dev \
    git \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Set up Hugging Face user (u:1000)
RUN useradd -m -u 1000 user
USER user
ENV PATH="/home/user/.local/bin:${PATH}"

WORKDIR /app

# Sync PyTorch
RUN pip install --no-cache-dir torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121

# Install Unsloth
RUN pip install --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

# Install requirements
COPY --chown=user requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Install Secondary (No-Deps)
RUN pip install --no-cache-dir --no-deps xformers<0.28 trl<0.13.0 peft accelerate bitsandbytes

# Copy training files
COPY --chown=user hf_training/ ./hf_training/
COPY --chown=user aegis_env/ ./aegis_env/
COPY --chown=user scripts/ ./scripts/
COPY --chown=user aegis_training_data_500.json .

# CMD
CMD ["python3", "hf_training/train.py"]