Bootstrap: docker
From: ollama/ollama:latest

%post
    # Prevent interactive prompts
    export DEBIAN_FRONTEND=noninteractive
    export TZ=UTC
    echo "Checking OS release"
    cat /etc/os-release || true
    echo "Installing system dependencies (Ubuntu/Debian)"
    # ollama/ollama base image is Ubuntu-based and includes CUDA runtime
    apt-get -y update
    apt-get -y install \
        python3 \
        python3-pip \
        python3-venv \
        git \
        curl \
        netcat-traditional \
        htop \
        openssh-client

    # Create application structure with proper permissions
    mkdir -p /app/{models,data,logs}
    chmod -R 777 /app
    mkdir -p /var/log/ollama
    chmod -R 777 /var/log/ollama

    # Create virtual environment
    python3 -m venv /app/venv
    # POSIX-compliant activation under /bin/sh
    . /app/venv/bin/activate

    echo "Checking NVIDIA GPU"
    # nvidia-smi
    
    # Upgrade pip and install Python dependencies compatible with CUDA 12.x and Ollama
    python -m pip install --upgrade pip

    # Core Python dependencies (latest versions)
    pip install --no-cache-dir \
        numpy==2.2.6 \
        pandas==3.0.0 \
        pyyaml==6.0.3 \
        requests==2.32.5 \
        tqdm==4.67.1 \
        psutil==7.2.2

    # Install PyTorch 2.9.1 with CUDA 12.9 wheels (pinned to match vLLM requirements)
    pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cu129 \
        torch==2.9.1 \
        torchvision==0.24.1 \
        torchaudio==2.9.1
    
    # Install vLLM for high-performance LLM inference
    # Pinned versions ensure PyTorch 2.9.1+cu129 is not overwritten
    pip install --no-cache-dir vllm==0.15.0

%environment
    # Locale
    export LC_ALL=C

    # Python virtual environment and CUDA paths
    export PATH=/app/venv/bin:/usr/local/cuda/bin:$PATH
    export VIRTUAL_ENV=/app/venv
    export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

%labels
    Author Rohan Marwaha, Darren Adams
    Version v2.0
    Description LLM Batch Processing Container - Simplified for SLURM
    LastModified 2025-10-24
    ModifiedBy Darren Adams
    Note Runtime configuration handled by SLURM job script