From bc0e12819d4af9db4f1331830ab91e685a13186c Mon Sep 17 00:00:00 2001
From: clsferguson <48876201+clsferguson@users.noreply.github.com>
Date: Sun, 21 Sep 2025 22:42:46 -0600
Subject: [PATCH] build(docker): compile SageAttention 2.2 on slim trixie with
 Debian CUDA toolkit; install wheel into runtime
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Switch to a two-stage build that uses python:3.12-slim-trixie as both builder and runtime, enabling contrib/non-free/non-free-firmware in APT to install Debian’s nvidia-cuda-toolkit (nvcc) for compiling SageAttention 2.2 from source. Install Torch cu129 in the builder and build a cp312 wheel, then copy and install that wheel into the slim runtime so --use-sage-attention works at startup. This removes the heavy CUDA devel base, avoids a venv by permitting pip system installs during build, and keeps the final image minimal while ensuring ABI alignment with Torch cu129.
---
 Dockerfile | 55 +++++++++++++++++++++++++++---------------------------
 1 file changed, 28 insertions(+), 27 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 16a7cc22f..2b4d4bd2b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,32 +1,33 @@
 # --------------------------
-# Stage 1: build SageAttention 2.2 wheel from source with nvcc available
+# Stage 1: build SageAttention 2.2 wheel (Debian trixie + nvcc)
 # --------------------------
-FROM nvidia/cuda:12.9.0-devel-ubuntu24.04 AS sage-builder
+FROM python:3.12.11-slim-trixie AS sage-builder
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_DISABLE_PIP_VERSION_CHECK=1 \
-    PIP_NO_CACHE_DIR=1 \
-    VENV=/opt/venv
+    PIP_NO_CACHE_DIR=1
 
-# Python 3.12 and build tools
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3 python3-pip python3-venv \
-    git build-essential cmake \
- && rm -rf /var/lib/apt/lists/*
-
-# Create a venv to avoid PEP 668 'externally-managed-environment'
-RUN python3 -m venv "$VENV"
-ENV PATH="$VENV/bin:$PATH"
+# Enable contrib/non-free/non-free-firmware, install CUDA toolkit + build deps
+RUN set -eux; \
+  sed -i 's/ main$/ main contrib non-free non-free-firmware/g' /etc/apt/sources.list; \
+  apt-get update; \
+  apt-get install -y --no-install-recommends \
+    ca-certificates curl git build-essential cmake \
+    nvidia-cuda-toolkit \
+  ; \
+  rm -rf /var/lib/apt/lists/*
 
 WORKDIR /tmp/sage
 
-# Install Torch (cu129) in the venv before building the extension so ABIs align
-RUN python -m pip install --upgrade pip setuptools wheel \
- && python -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129
+# Install Torch cu129 in system site (PEP 668: allow system installs)
+RUN python -m pip install --upgrade pip setuptools wheel --break-system-packages && \
+    python -m pip install torch torchvision torchaudio \
+      --extra-index-url https://download.pytorch.org/whl/cu129 \
+      --break-system-packages
 
-# Shallow clone latest SageAttention and build a cp312 wheel from source
-RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git . \
- && python -m pip wheel . --no-deps --no-build-isolation -w /dist
+# Shallow clone SageAttention and build a cp312 wheel
+RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git . && \
+    python -m pip wheel . --no-deps --no-build-isolation -w /dist --break-system-packages
 
 # --------------------------
 # Stage 2: runtime image (slim)
@@ -41,8 +42,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 # System deps
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    git build-essential cmake \
-    libgl1 libglx-mesa0 libglib2.0-0 \
+    git build-essential cmake libgl1 libglx-mesa0 libglib2.0-0 \
     fonts-dejavu-core fontconfig util-linux \
  && rm -rf /var/lib/apt/lists/*
 
@@ -52,12 +52,13 @@ RUN groupadd --gid 1000 appuser \
 
 WORKDIR /app/ComfyUI
 
-# Install core deps
+# Install core deps (Torch cu129 must match builder)
 COPY requirements.txt ./
-RUN python -m pip install --upgrade pip setuptools wheel \
- && python -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 \
- && python -m pip install -r requirements.txt \
- && python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py
+RUN python -m pip install --upgrade pip setuptools wheel && \
+    python -m pip install torch torchvision torchaudio \
+      --extra-index-url https://download.pytorch.org/whl/cu129 && \
+    python -m pip install -r requirements.txt && \
+    python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py
 
 # Install the SageAttention wheel built in the builder stage
 COPY --from=sage-builder /dist/sageattention-*.whl /tmp/
@@ -66,7 +67,7 @@ RUN python -m pip install /tmp/sageattention-*.whl
 # Copy the application
 COPY . .
 
-# Entrypoint
+# Entrypoint and launch
 COPY entrypoint.sh /entrypoint.sh
 RUN chmod +x /entrypoint.sh \
  && chown -R appuser:appuser /app /home/appuser /entrypoint.sh