FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-devel

# Set environment variables
ENV PYTHONUNBUFFERED=1
ENV HF_HOME=/app/.cache/huggingface
ENV TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers
ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    curl \
    wget \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Download orbital CLI from HuggingFace (has native validation, no kflow needed)
RUN wget -q https://huggingface.co/javasop/orbital-cli/resolve/main/orbital-linux-x86_64 -O /usr/local/bin/orbital \
    && chmod +x /usr/local/bin/orbital \
    && orbital --version || echo "Orbital CLI installed"

# Copy requirements and install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt

# Copy training code
COPY . .

# Create cache directories
RUN mkdir -p /app/.cache/huggingface

# Default command - uses HuggingFace config
CMD ["python", "train.py", "--config", "config-huggingface.yaml"]