Dockerfile.runpod

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

ENV DEBIAN_FRONTEND=noninteractive
ENV TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0"
ENV PYTHONUNBUFFERED=1

# HuggingFace settings
ENV HF_HUB_ENABLE_HF_TRANSFER=1
ENV HF_HOME=/runpod-volume/huggingface-cache
ENV HUGGINGFACE_HUB_CACHE=/runpod-volume/huggingface-cache/hub
ENV TRANSFORMERS_CACHE=/runpod-volume/huggingface-cache/hub

# Disable telemetry
ENV NO_ALBUMENTATIONS_UPDATE=1
ENV DISABLE_TELEMETRY=YES

# Install system deps
RUN apt-get update && apt-get install -y --no-install-recommends \
    git git-lfs curl wget python3.10 python3.10-dev python3-pip \
    ffmpeg libgl1-mesa-glx libglib2.0-0 aria2 \
    && rm -rf /var/lib/apt/lists/*

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1
RUN pip install --upgrade pip

# Install 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/cu124

# Install RunPod and HF transfer
RUN pip install --no-cache-dir runpod hf_transfer huggingface_hub

# Copy project
WORKDIR /app
COPY ai-toolkit /app/ai-toolkit
COPY rp_handler.py /app/rp_handler.py

# Install ai-toolkit requirements
RUN pip install --no-cache-dir -r /app/ai-toolkit/requirements.txt

# Create workspace directories
RUN mkdir -p /workspace/dataset /workspace/output

WORKDIR /app
CMD ["python", "-u", "rp_handler.py"]