Add Dockerfile to examples/puzzletron

examples/puzzletron/README.md

@@ -43,7 +43,45 @@ extension against that same installation; mixing PyTorch or CUDA builds can
 cause import failures or incorrect GPU execution. AIPerf uses the official PyPI
 package; no custom AIPerf fork is required.
 
-### 1. Start from the CUDA development image
+### Quick start
+
+Two automated options are provided. Both are driven by the same script so
+there is no duplication between them:
+
+**Option A — Docker image** (recommended for reproducibility):
+
+```bash
+# Build from the repo root so the full source tree is available as build context
+docker build \
+  -f examples/puzzletron/configs/families/nemotron3/Dockerfile \
+  -t puzzletron-nemotron3 \
+  .
+```
+
+**Option B — Bare-metal / cluster** (no Docker required):
+
+```bash
+export MODEL_OPT_ROOT=/path/to/modelopt   # defaults to repo root when run from there
+export VLLM_ROOT=/workspace/vllm
+export AUTOMODEL_ROOT=/workspace/Automodel
+export VIRTUAL_ENV=/venv
+
+# SKIP_APT=1 if build-essential/cmake/git/ninja are already installed
+bash examples/puzzletron/configs/families/nemotron3/setup_env.sh
+
+source "${VIRTUAL_ENV}/bin/activate"
+```
+
+The script is idempotent: existing git clones and the venv are reused.
+`FORCE_CUDA=1` and `TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0"` are set by
+default so CUDA extensions build without a live GPU.
+
+### Manual setup (reference)
+
+The steps below document what the script does. Follow them if you need to
+customise individual stages or debug a failed install.
+
+#### 1. Start from the CUDA development image
 
 Use this image for local containers and cluster jobs:
 
@@ -71,7 +109,7 @@ DEBIAN_FRONTEND=noninteractive apt-get install -y \
   python3 python3-dev python3-pip python3-venv
 ```
 
-### 2. Clone the tracked forks
+#### 2. Clone the tracked forks
 
 Keep ModelOpt and the two Puzzletron forks as siblings:
 
@@ -93,7 +131,7 @@ git clone --branch puzzletron --single-branch \
 └── Automodel/
 ```
 
-### 3. Create the environment and install runtime packages
+#### 3. Create the environment and install runtime packages
 
 The patched vLLM branch uses PyTorch 2.11.0 with CUDA 12.9. Install that
 combination before anything that compiles CUDA code:
@@ -116,7 +154,6 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \
 python -m pip install -e "${AUTOMODEL_ROOT}"
 python -m pip install aiperf
 python -m pip install -e "${MODEL_OPT_ROOT}[hf]"
-python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements.txt"
 ```
 
 Do not add `--no-deps`: the packages need their declared Python dependencies.
@@ -127,7 +164,8 @@ Install only the model-specific kernels required by the target architecture:
 
 ```bash
 # Mixture of experts
-python -m pip install --no-build-isolation \
+FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0" \
+  python -m pip install --no-build-isolation \
   "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4"
 
 # Mamba
@@ -137,7 +175,7 @@ python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation
 python -m pip install "flash-linear-attention[cuda]"
 ```
 
-### 4. Verify the exact environment
+#### 4. Verify the exact environment
 
 Run these checks inside the same container and venv used for Puzzletron jobs:
 

examples/puzzletron/configs/families/nemotron3/Dockerfile

@@ -0,0 +1,32 @@
+ARG CUDA_VERSION=12.9.2
+ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04
+
+FROM ${BUILD_BASE_IMAGE} AS base
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ARG MODEL_OPT_ROOT=/workspace/modelopt
+ARG VLLM_ROOT=/workspace/vllm
+ARG AUTOMODEL_ROOT=/workspace/Automodel
+
+# Copy only the setup script first so that the expensive dep-install layer
+# stays cached as long as the script hasn't changed, regardless of source edits.
+COPY examples/puzzletron/configs/families/nemotron3/setup_env.sh /tmp/setup_env.sh
+
+# Tell the build system to proceed without a live GPU
+ARG FORCE_CUDA=1
+# Avoid GPU detection at compile time
+ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0"
+
+RUN MODEL_OPT_ROOT="${MODEL_OPT_ROOT}" \
+    VLLM_ROOT="${VLLM_ROOT}" \
+    AUTOMODEL_ROOT="${AUTOMODEL_ROOT}" \
+    bash /tmp/setup_env.sh --deps
+
+ENV VIRTUAL_ENV=/venv
+ENV PATH=/venv/bin/:$PATH
+
+# Copy the full repo after deps so source changes don't bust the cache above.
+COPY . "${MODEL_OPT_ROOT}"/
+
+RUN MODEL_OPT_ROOT="${MODEL_OPT_ROOT}" bash /tmp/setup_env.sh --modelopt

examples/puzzletron/configs/families/nemotron3/setup_env.sh

@@ -0,0 +1,119 @@
+#!/usr/bin/env bash
+# Shared setup script for the Nemotron3 Puzzletron environment.
+#
+# Used by the Dockerfile (--deps before COPY, --modelopt after COPY) and
+# can be run standalone to set up a bare-metal environment (no arguments).
+#
+# Usage:
+#   ./setup_env.sh [--deps | --modelopt]
+#     --deps      System packages, git clones, venv, pip deps — safe to Docker-cache
+#     --modelopt  Install the local ModelOpt source only (run after --deps)
+#     (no args)   Full setup — suitable for bare-metal
+#
+# Configurable via environment variables (all have defaults):
+#   MODEL_OPT_ROOT         Path to the ModelOpt repo root
+#                          Default: 5 directories above this script
+#   VLLM_ROOT              Where to clone vLLM         (default: /workspace/vllm)
+#   AUTOMODEL_ROOT         Where to clone Automodel    (default: /workspace/Automodel)
+#   VIRTUAL_ENV            Python venv path            (default: /venv)
+#   SKIP_APT               Set to 1 to skip apt-get   (default: 0)
+#   FORCE_CUDA             Build CUDA exts without a live GPU (default: 1)
+#   TORCH_CUDA_ARCH_LIST   Architectures to compile for
+#                          (default: "8.0;8.6;9.0;10.0")
+
+set -euo pipefail
+
+MODE="${1:-all}"
+
+# ── Paths ──────────────────────────────────────────────────────────────────────
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# examples/puzzletron/configs/families/nemotron3 → 5 levels up to repo root
+MODEL_OPT_ROOT="${MODEL_OPT_ROOT:-$(cd "${SCRIPT_DIR}/../../../../.." && pwd)}"
+VLLM_ROOT="${VLLM_ROOT:-/workspace/vllm}"
+AUTOMODEL_ROOT="${AUTOMODEL_ROOT:-/workspace/Automodel}"
+VIRTUAL_ENV="${VIRTUAL_ENV:-/venv}"
+
+# ── Build flags ────────────────────────────────────────────────────────────────
+SKIP_APT="${SKIP_APT:-0}"
+export FORCE_CUDA="${FORCE_CUDA:-1}"
+export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.0;8.6;9.0;10.0}"
+
+activate_venv() {
+    export VIRTUAL_ENV
+    export PATH="${VIRTUAL_ENV}/bin:${PATH}"
+}
+
+# ══════════════════════════════════════════════════════════════════════════════
+# DEPS phase — nothing here depends on the local ModelOpt source.
+# In Docker this runs before `COPY .` so it stays cached across source changes.
+# ══════════════════════════════════════════════════════════════════════════════
+install_deps() {
+    if [[ "${SKIP_APT}" != "1" ]]; then
+        apt-get update
+        apt-get install -y build-essential cmake git ninja-build \
+            python3 python3-dev python3-pip python3-venv
+    fi
+
+    # Clone external repos (idempotent)
+    [[ -d "${VLLM_ROOT}" ]] || \
+        git clone --branch feature/add_anymodel_to_vllm --single-branch \
+            https://github.com/Separius/vllm.git "${VLLM_ROOT}"
+    [[ -d "${AUTOMODEL_ROOT}" ]] || \
+        git clone --branch puzzletron --single-branch \
+            https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}"
+
+    [[ -d "${VIRTUAL_ENV}" ]] || python3 -m venv "${VIRTUAL_ENV}"
+    activate_venv
+
+    python -m pip install --upgrade pip \
+        "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust wheel \
+        "packaging>=24.2" "cmake>=3.26.1" ninja jinja2
+
+    python -m pip install \
+        torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 \
+        --index-url https://download.pytorch.org/whl/cu129
+
+    VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \
+        python -m pip install --no-build-isolation -e "${VLLM_ROOT}"
+
+    python -m pip install -e "${AUTOMODEL_ROOT}"
+    python -m pip install aiperf
+
+    # CUDA extension packages placed here (before COPY .) so they stay cached
+    # across ModelOpt source changes in Docker builds.
+    python -m pip install --no-build-isolation \
+        "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4"
+    python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation
+    python -m pip install "flash-linear-attention[cuda]"
+}
+
+# ══════════════════════════════════════════════════════════════════════════════
+# MODELOPT phase — installs the local ModelOpt source.
+# In Docker this runs after `COPY .` and re-runs on every source change.
+# ══════════════════════════════════════════════════════════════════════════════
+install_modelopt() {
+    activate_venv
+    python -m pip install -e "${MODEL_OPT_ROOT}[hf]"
+}
+
+# ── Dispatch ───────────────────────────────────────────────────────────────────
+case "${MODE}" in
+    --deps)
+        install_deps
+        ;;
+    --modelopt)
+        install_modelopt
+        ;;
+    all)
+        install_deps
+        install_modelopt
+        echo ""
+        echo "Setup complete. Activate your environment with:"
+        echo "  source ${VIRTUAL_ENV}/bin/activate"
+        ;;
+    *)
+        echo "Unknown argument: ${MODE}" >&2
+        echo "Usage: $0 [--deps | --modelopt]" >&2
+        exit 1
+        ;;
+esac