WIP: migrate base image to NGC PyTorch (Blackwell/CUDA-13) — separate from DR-1170 fix

.github/workflows/manual-build-all.yml

@@ -60,57 +60,6 @@ jobs:
             *.args.RELEASE_VERSION=${{ env.RELEASE_VERSION }}
             *.args.HUGGINGFACE_ACCESS_TOKEN=${{ env.HUGGINGFACE_ACCESS_TOKEN }}
 
-  build-base-cuda12-8-1:
-    runs-on: [blacksmith-8vcpu-ubuntu-2204, linux]
-    permissions:
-      contents: read
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-
-      - name: Clear space to remove unused folders
-        run: |
-          rm -rf /usr/share/dotnet
-          rm -rf /opt/ghc
-          rm -rf "/usr/local/share/boost"
-          rm -rf "$AGENT_TOOLSDIRECTORY"
-
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v3
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Login to Docker Hub
-        uses: docker/login-action@v3
-        with:
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-
-      - name: blacksmith docker layer cache
-        uses: useblacksmith/build-push-action@v1
-        with:
-          setup-only: true
-
-      - name: Set environment variables
-        run: |
-          echo "DOCKERHUB_REPO=${{ vars.DOCKERHUB_REPO }}" >> $GITHUB_ENV
-          echo "DOCKERHUB_IMG=${{ vars.DOCKERHUB_IMG }}" >> $GITHUB_ENV
-          echo "HUGGINGFACE_ACCESS_TOKEN=${{ secrets.HUGGINGFACE_ACCESS_TOKEN }}" >> $GITHUB_ENV
-          echo "RELEASE_VERSION=${{ github.event.inputs.version }}" >> $GITHUB_ENV
-
-      - name: Build and push base-cuda12-8-1 image
-        uses: docker/bake-action@v2
-        with:
-          push: true
-          targets: base-cuda12-8-1
-          set: |
-            *.args.DOCKERHUB_REPO=${{ env.DOCKERHUB_REPO }}
-            *.args.DOCKERHUB_IMG=${{ env.DOCKERHUB_IMG }}
-            *.args.RELEASE_VERSION=${{ env.RELEASE_VERSION }}
-            *.args.HUGGINGFACE_ACCESS_TOKEN=${{ env.HUGGINGFACE_ACCESS_TOKEN }}
-
   build-sdxl:
     runs-on: [blacksmith-8vcpu-ubuntu-2204, linux]
     permissions:
@@ -422,7 +371,6 @@ jobs:
     needs:
       [
         build-base,
-        build-base-cuda12-8-1,
         build-sdxl,
         build-sd3,
         build-flux1-schnell,

.github/workflows/release.yml

@@ -85,7 +85,6 @@ jobs:
           - flux1-dev
           - flux1-dev-fp8
           - z-image-turbo
-          - base-cuda12-8-1
 
     steps:
       - name: Checkout

Dockerfile

@@ -1,5 +1,6 @@
-# Build argument for base image selection
-ARG BASE_IMAGE=nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04
+# Build argument for base image selection. Default: NVIDIA NGC PyTorch container
+# (Ubuntu 24.04 + tuned torch/cuDNN/NCCL incl. Blackwell sm_120 kernels).
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.05-py3
 
 # Stage 1: Base image with common dependencies
 FROM ${BASE_IMAGE} AS base
@@ -9,6 +10,10 @@ ARG COMFYUI_VERSION=latest
 ARG CUDA_VERSION_FOR_COMFY
 ARG ENABLE_PYTORCH_UPGRADE=false
 ARG PYTORCH_INDEX_URL
+# When the base image already ships a tuned PyTorch (e.g. NGC nvcr.io/nvidia/pytorch),
+# keep "true" to reuse it instead of letting comfy-cli install its own wheel.
+# Set to "false" only when using a plain CUDA base that has no torch.
+ARG BASE_PROVIDES_TORCH=true
 
 # Prevents prompts from packages asking for user input during installation
 ENV DEBIAN_FRONTEND=noninteractive
@@ -38,11 +43,17 @@ RUN apt-get update && apt-get install -y \
 # Clean up to reduce image size
 RUN apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/*
 
-# Install uv (latest) using official installer and create isolated venv
+# Install uv (latest) using official installer and create the venv.
+# With BASE_PROVIDES_TORCH=true the venv inherits the base image's site-packages
+# (--system-site-packages) so NVIDIA's bundled torch/cuDNN/NCCL stay visible.
 RUN wget -qO- https://astral.sh/uv/install.sh | sh \
     && ln -s /root/.local/bin/uv /usr/local/bin/uv \
     && ln -s /root/.local/bin/uvx /usr/local/bin/uvx \
-    && uv venv /opt/venv
+    && if [ "$BASE_PROVIDES_TORCH" = "true" ]; then \
+         uv venv --system-site-packages --python /usr/bin/python3.12 /opt/venv; \
+       else \
+         uv venv /opt/venv; \
+       fi
 
 # Use the virtual environment for all subsequent commands
 ENV PATH="/opt/venv/bin:${PATH}"
@@ -51,7 +62,11 @@ ENV PATH="/opt/venv/bin:${PATH}"
 RUN uv pip install comfy-cli pip setuptools wheel
 
 # Install ComfyUI
-RUN if [ -n "${CUDA_VERSION_FOR_COMFY}" ]; then \
+# When the base image provides torch, skip comfy-cli's torch install so it does
+# not clobber the bundled wheel; otherwise install torch for the requested CUDA.
+RUN if [ "$BASE_PROVIDES_TORCH" = "true" ]; then \
+      /usr/bin/yes | comfy --workspace /comfyui install --version "${COMFYUI_VERSION}" --skip-torch-or-directml --nvidia; \
+    elif [ -n "${CUDA_VERSION_FOR_COMFY}" ]; then \
       /usr/bin/yes | comfy --workspace /comfyui install --version "${COMFYUI_VERSION}" --cuda-version "${CUDA_VERSION_FOR_COMFY}" --nvidia; \
     else \
       /usr/bin/yes | comfy --workspace /comfyui install --version "${COMFYUI_VERSION}" --nvidia; \
@@ -62,6 +77,30 @@ RUN if [ "$ENABLE_PYTORCH_UPGRADE" = "true" ]; then \
       uv pip install --force-reinstall torch torchvision torchaudio --index-url ${PYTORCH_INDEX_URL}; \
     fi
 
+# comfy-cli installs ComfyUI into its OWN workspace venv (/comfyui/.venv), but
+# start.sh launches ComfyUI with /opt/venv's python. That mismatch means the
+# launch venv is missing ComfyUI's runtime deps (e.g. sqlalchemy, pulled in by
+# ComfyUI's asset DB), so ComfyUI crashes at startup and surfaces as the
+# misleading "ComfyUI server not reachable" error. Mirror ComfyUI's full
+# dependency set (core + custom nodes) into /opt/venv so the launch venv is
+# complete. This is the root-cause fix for DR-1170.
+RUN uv pip install -r /comfyui/requirements.txt \
+    && for r in /comfyui/custom_nodes/*/requirements.txt; do \
+         [ -f "$r" ] && uv pip install -r "$r" || true; \
+       done
+
+# Pin ComfyUI's unbounded ML dependencies to their last known-good majors.
+# ComfyUI requires transformers>=4.50.3 and huggingface-hub with NO upper bound,
+# so a fresh install pulls transformers 5.x + huggingface-hub 1.x whose breaking
+# API changes also break ComfyUI startup. Keep them on the last good major.
+RUN uv pip install "transformers>=4.50.3,<5" "huggingface-hub<1.0"
+
+# Build-time smoke test: actually start ComfyUI (imports the full node graph) so
+# a startup-breaking dependency is caught HERE, at build time, instead of as a
+# runtime "server not reachable" failure on a serverless worker. Runs on CPU —
+# no GPU is needed to exercise the import graph.
+RUN cd /comfyui && timeout 300 python main.py --quick-test-for-ci --cpu
+
 # Change working directory to ComfyUI
 WORKDIR /comfyui
 

docker-bake.hcl

@@ -14,13 +14,24 @@ variable "COMFYUI_VERSION" {
   default = "latest"
 }
 
-# Global defaults for standard CUDA 12.6.3 images
+# Default base image: NVIDIA NGC PyTorch container (Ubuntu 24.04 + tuned
+# torch/cuDNN/NCCL with Blackwell sm_120 kernels + cuda-compat forward-compat
+# libs). Validated end-to-end on every serverless GPU (Ampere → Blackwell).
 variable "BASE_IMAGE" {
-  default = "nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04"
+  default = "nvcr.io/nvidia/pytorch:26.05-py3"
 }
 
+# The NGC base already ships a tuned PyTorch — reuse it (venv built with
+# --system-site-packages, comfy-cli run with --skip-torch-or-directml) instead
+# of letting comfy-cli install its own wheel.
+variable "BASE_PROVIDES_TORCH" {
+  default = "true"
+}
+
+# Empty: do not have comfy-cli install torch for a specific CUDA (the base
+# provides it). Only used when BASE_PROVIDES_TORCH=false.
 variable "CUDA_VERSION_FOR_COMFY" {
-  default = "12.6"
+  default = ""
 }
 
 variable "ENABLE_PYTORCH_UPGRADE" {
@@ -36,7 +47,7 @@ variable "HUGGINGFACE_ACCESS_TOKEN" {
 }
 
 group "default" {
-  targets = ["base", "sdxl", "sd3", "flux1-schnell", "flux1-dev", "flux1-dev-fp8", "z-image-turbo", "base-cuda12-8-1"]
+  targets = ["base", "sdxl", "sd3", "flux1-schnell", "flux1-dev", "flux1-dev-fp8", "z-image-turbo"]
 }
 
 target "base" {
@@ -46,6 +57,7 @@ target "base" {
   platforms = ["linux/amd64"]
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -61,6 +73,7 @@ target "sdxl" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -77,6 +90,7 @@ target "sd3" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -94,6 +108,7 @@ target "flux1-schnell" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -111,6 +126,7 @@ target "flux1-dev" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -128,6 +144,7 @@ target "flux1-dev-fp8" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -144,6 +161,7 @@ target "z-image-turbo" {
   target = "final"
   args = {
     BASE_IMAGE = "${BASE_IMAGE}"
+    BASE_PROVIDES_TORCH = "${BASE_PROVIDES_TORCH}"
     COMFYUI_VERSION = "${COMFYUI_VERSION}"
     CUDA_VERSION_FOR_COMFY = "${CUDA_VERSION_FOR_COMFY}"
     ENABLE_PYTORCH_UPGRADE = "${ENABLE_PYTORCH_UPGRADE}"
@@ -154,20 +172,3 @@ target "z-image-turbo" {
   tags = ["${DOCKERHUB_REPO}/${DOCKERHUB_IMG}:${RELEASE_VERSION}-z-image-turbo"]
   inherits = ["base"]
 }
-
-target "base-cuda12-8-1" {
-  context = "."
-  dockerfile = "Dockerfile"
-  target = "base"
-  platforms = ["linux/amd64"]
-  args = {
-    BASE_IMAGE = "nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04"
-    COMFYUI_VERSION = "${COMFYUI_VERSION}"
-    CUDA_VERSION_FOR_COMFY = ""
-    ENABLE_PYTORCH_UPGRADE = "true"
-    PYTORCH_INDEX_URL = "https://download.pytorch.org/whl/cu128"
-    MODEL_TYPE = "base"
-  }
-  tags = ["${DOCKERHUB_REPO}/${DOCKERHUB_IMG}:${RELEASE_VERSION}-base-cuda12.8.1"]
-}
-

handler.py

@@ -26,6 +26,53 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# ---------------------------------------------------------------------------
+# Optional host/GPU diagnostics. When REPORT_HOST_CUDA=true, every job result
+# includes a "host" block with the host's actual CUDA version (from nvidia-smi),
+# driver, GPU name, and the torch build CUDA. Used to validate which exact CUDA
+# version a worker landed on (the serverless API only exposes the min-cuda floor).
+# Cached after first call so it does not run per job.
+# ---------------------------------------------------------------------------
+_HOST_INFO_CACHE = None
+
+
+def _get_host_info():
+    global _HOST_INFO_CACHE
+    if _HOST_INFO_CACHE is not None:
+        return _HOST_INFO_CACHE
+    import subprocess
+
+    info = {}
+    try:
+        out = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader"],
+            capture_output=True, text=True, timeout=15,
+        ).stdout.strip().splitlines()
+        if out:
+            name, driver = (out[0].split(",") + [""])[:2]
+            info["gpu"] = name.strip()
+            info["driver"] = driver.strip()
+        # nvidia-smi header reports the max CUDA version the driver supports
+        hdr = subprocess.run(["nvidia-smi"], capture_output=True, text=True, timeout=15).stdout
+        for line in hdr.splitlines():
+            if "CUDA Version" in line:
+                info["host_cuda"] = line.split("CUDA Version:")[1].split()[0].strip()
+                break
+    except Exception as exc:
+        info["nvidia_smi_error"] = str(exc)
+    try:
+        import torch
+
+        info["torch"] = torch.__version__
+        info["torch_cuda"] = torch.version.cuda
+        if torch.cuda.is_available():
+            cap = torch.cuda.get_device_capability(0)
+            info["sm"] = f"sm_{cap[0]}{cap[1]}"
+    except Exception as exc:
+        info["torch_error"] = str(exc)
+    _HOST_INFO_CACHE = info
+    return info
+
 # Time to wait between API check attempts in milliseconds
 COMFY_API_AVAILABLE_INTERVAL_MS = int(
     os.environ.get("COMFY_API_AVAILABLE_INTERVAL_MS", 50)
@@ -872,6 +919,9 @@ def handler(job):
 
     final_result = {}
 
+    if os.environ.get("REPORT_HOST_CUDA", "false").lower() == "true":
+        final_result["host"] = _get_host_info()
+
     if output_data:
         final_result["images"] = output_data
 

src/start.sh

@@ -34,15 +34,25 @@ import torch
 try:
     torch.cuda.init()
     name = torch.cuda.get_device_name(0)
-    print(f'OK: {name}')
+    cap = torch.cuda.get_device_capability(0)
+    # Launch a real kernel. The driver-only calls above succeed even when this
+    # PyTorch build has no compiled kernels for the GPU architecture (e.g. cu126
+    # wheels on a Blackwell sm_120 device). Without this, the worker boots,
+    # ComfyUI dies on the first GPU op, and it surfaces as the misleading
+    # 'ComfyUI server not reachable' error instead of a clear cause here.
+    _ = (torch.zeros(8, device='cuda') + 1).sum().item()
+    torch.cuda.synchronize()
+    print(f'OK: {name} (sm_{cap[0]}{cap[1]}), torch {torch.__version__}, cuda {torch.version.cuda}')
 except Exception as e:
     print(f'FAIL: {e}')
     exit(1)
 " 2>&1); then
-    echo "worker-comfyui: GPU is not available. PyTorch CUDA init failed:"
+    echo "worker-comfyui: GPU is not available or incompatible with this PyTorch build:"
     echo "worker-comfyui: $GPU_CHECK"
-    echo "worker-comfyui: This usually means the GPU on this machine is not properly initialized."
-    echo "worker-comfyui: Please contact RunPod support and report this machine."
+    echo "worker-comfyui: A 'no kernel image is available' error means this torch build"
+    echo "worker-comfyui: lacks kernels for this GPU (e.g. cu126 on a Blackwell sm_120 card)."
+    echo "worker-comfyui: Otherwise the GPU may not be properly initialized — please"
+    echo "worker-comfyui: contact RunPod support and report this machine."
     exit 1
 fi
 echo "worker-comfyui: GPU available — $GPU_CHECK"