Bootstrap: docker
From: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

%environment
    export LANG=C.UTF-8
    export LC_ALL=C.UTF-8
    export CUTLASS_PATH=/kernels/cutlass
    export LAYERNORM_TYPE=fast_layernorm
    export PATH=/usr/local/bin:$PATH

%post

    # Install basic packages for compiling and building
    export DEBIAN_FRONTEND=noninteractive
    apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \
        build-essential \
        cmake \
        g++-12 \
        git \
        curl \
        wget \
        ca-certificates \
        libjpeg-dev \
        libpng-dev \
        librdmacm1 \
        libibverbs1 \
        ibverbs-providers \
        tzdata \
        libgl1-mesa-glx \
        libglib2.0-0 \
        fontconfig \
        && rm -rf /var/lib/apt/lists/*
    
    # Install Miniconda & use Python 3.10
    export PYTHON_VERSION=3.10
    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/install-conda.sh
    chmod +x /tmp/install-conda.sh
    bash /tmp/install-conda.sh -b -f -p /usr/local
    rm -f /tmp/install-conda.sh
    
    # Initialize conda
    export PATH=/usr/local/bin:$PATH
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
    conda install -y python=${PYTHON_VERSION}
    
    # Install packages from requirements.txt with pip
    # Note: Copy files section will handle copying the project
    cd /app/ODesign
    pip install --no-cache-dir -r requirements.txt -f https://data.pyg.org/whl/torch-2.3.1+cu121.html
    
    # Set up kernels
    # DS4Sci_EvoformerAttention
    mkdir -p /kernels
    git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git /kernels/cutlass

%files
    . /app/ODesign

%labels
    Author ODesign
    Version 1.0