Dockerfile

FROM nvidia/cuda:12.9.0-cudnn-devel-ubuntu24.04

# Install system dependencies including Python
RUN apt-get update && apt-get install -y \
    python3 \
    python3-pip \
    python3-venv \
    python3-full \
    ninja-build \
    && rm -rf /var/lib/apt/lists/*

# Set CUDA architecture to avoid compiling for all architectures
# For Ampere GPUs (A100)
ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
ENV XLSTM_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'
# Fix CUDA template linking issue by enabling relocatable device code
ENV CUDA_NVCC_FLAGS="-rdc=true --expt-relaxed-constexpr"
# Limit parallel compilation jobs
ENV MAX_JOBS=4
# Set CUDA Home
ENV CUDA_HOME="/usr/local/cuda"
# Set working directory
WORKDIR /app

# Copy the entire project first (including submodules)
COPY . .

# Create a virtual environment and install dependencies
RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# Install dependencies using pip
RUN pip install --upgrade pip
# Install PyTorch compatible with CUDA 12.9
RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129

# Install requirements
RUN pip install -r requirements.txt

# Install Runpod SDK
RUN pip install --no-cache-dir runpod

# Start the container
CMD ["python3", "rp_handler.py"]