Dockerfile

# ==============================================================================
# AB-GCG Dockerfile for RunPod (4×A100 80GB)
# 
# Base: PyTorch 2.4.0 + CUDA 12.1 (RunPod compatible)
# Key constraints:
#   - nanogcg requires transformers>=4.4,<=4.47.1
#   - Qwen-2.5-3B-Instruct benefits from flash-attn 2
#   - analysis scripts need matplotlib
#
# Build:
#   docker build -t ab-gcg:latest .
#
# Run (RunPod or local):
#   docker run --gpus all -it --shm-size=16g \
#     -v /workspace:/workspace \
#     ab-gcg:latest /bin/bash
# ==============================================================================

FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-devel

LABEL maintainer="AB-GCG Research"
LABEL description="AB-GCG: Adaptive Bandit-guided GCG attack experiments"

# Avoid interactive prompts during build
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV TORCH_CUDA_ARCH_LIST="8.0"

# ---- System dependencies ----
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    wget \
    curl \
    tmux \
    htop \
    nvtop \
    vim \
    unzip \
    && rm -rf /var/lib/apt/lists/*

# ---- Python dependencies ----
# Install transformers with nanogcg-compatible version ceiling
RUN pip install --no-cache-dir \
    "transformers>=4.44,<=4.47.1" \
    accelerate \
    sentencepiece \
    protobuf \
    scipy \
    numpy \
    tqdm

# Install nanogcg
RUN pip install --no-cache-dir nanogcg==0.3.0

# Install flash-attn 2 (speeds up attention computation on A100)
# Building from source; needs the devel image for nvcc
RUN pip install --no-cache-dir packaging ninja
RUN MAX_JOBS=4 pip install --no-cache-dir flash-attn --no-build-isolation

# Analysis & visualization dependencies
RUN pip install --no-cache-dir \
    matplotlib \
    pandas \
    seaborn

# ---- Project code is NOT baked in ----
# Code is cloned via git at runtime:
#   cd /workspace && git clone <repo_url> AB-GCG
# This keeps the image stable (only rebuild when dependencies change).
WORKDIR /workspace

# ---- Pre-download model (optional, uncomment to bake into image) ----
# This adds ~6GB to the image but eliminates download time at runtime.
# If you prefer to download at runtime, keep this commented out and
# the model will be cached in /root/.cache/huggingface on first run.
#
# RUN python -c "\
# from transformers import AutoModelForCausalLM, AutoTokenizer; \
# AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-3B-Instruct', torch_dtype='auto'); \
# AutoTokenizer.from_pretrained('Qwen/Qwen2.5-3B-Instruct')"

# ---- Default entrypoint ----
# Print welcome info, then sleep to keep the container alive.
# SSH in and run: cd /workspace/AB-GCG && bash scripts/start.sh
CMD ["/bin/bash", "-c", "echo '=== AB-GCG Environment Ready ===' && echo 'Model: Qwen/Qwen2.5-3B-Instruct' && echo '' && echo 'SSH in and run:' && echo '  cd /workspace/AB-GCG && bash scripts/start.sh' && echo '' && echo 'Or directly:' && echo '  python run_baseline.py / python run_ab_gcg.py' && sleep infinity"]