wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-15 13:02:35 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

feat(entrypoint): safer builds, dynamic CUDA archs, corrected git clone, first-run override, clarified Sage flags

Browse Source 
Cap build parallelism via MAX_JOBS (override SAGEATTENTION_MAX_JOBS) and
CMAKE_BUILD_PARALLEL_LEVEL to prevent OOM kills during nvcc/cc1plus when
ninja fanout is high in constrained containers.

Compute TORCH_CUDA_ARCH_LIST from torch.cuda device properties to target
exact GPU SMs across mixed setups; keep human-readable nvidia-smi logs.

Move PATH/PYTHONPATH exports earlier and use `python -m pip` with
`--no-cache-dir` consistently to avoid stale caches and reduce image bloat.

Fix git clone/update commands to standard HTTPS and reset against
origin/HEAD; keep shallow operations for speed and reproducibility.

Clarify Sage Attention flags: set SAGE_ATTENTION_AVAILABLE only when
module import succeeds; require FORCE_SAGE_ATTENTION=1 to enable at boot.

Keep first-run dependency installation with COMFY_AUTO_INSTALL=1 override
to re-run installs on later boots without removing the first-run flag.
This commit is contained in:
clsferguson 2025-09-26 18:19:23 -06:00 committed by GitHub
parent 30ed9ae7cf
commit 555b7d5606
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 24 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
entrypoint.sh
View File

@ -187,7 +187,7 @@ build_sage_attention_mixed() {
esac
log "Building Sage Attention..."
if MAX_JOBS=$(nproc) python -m pip install --no-cache-dir --user --no-build-isolation .; then
if MAX_JOBS="${MAX_JOBS:-2}" python -m pip install --no-cache-dir --user --no-build-isolation .; then
echo "$SAGE_STRATEGY" > "$SAGE_ATTENTION_BUILT_FLAG"
log "Sage Attention built successfully for strategy: $SAGE_STRATEGY"
cd "$BASE_DIR"
@ -237,26 +237,22 @@ except Exception as e:
# Main GPU detection and Sage Attention setup
setup_sage_attention() {
# Internal tracking and exported availability flag
export SAGE_ATTENTION_BUILT=0
export SAGE_ATTENTION_AVAILABLE=0
# Detect GPU generations
if ! detect_gpu_generations; then
log "No GPUs detected, skipping Sage Attention setup"
return 0
fi
# Determine optimal strategy
determine_sage_strategy
# Build/install if needed
if needs_rebuild || ! test_sage_attention; then
log "Building Sage Attention..."
if install_triton_version && build_sage_attention_mixed && test_sage_attention; then
export SAGE_ATTENTION_BUILT=1
export SAGE_ATTENTION_AVAILABLE=1
log "Sage Attention is built and importable; enable at boot by setting FORCE_SAGE_ATTENTION=1"
log "Sage Attention is built and importable; set FORCE_SAGE_ATTENTION=1 to enable at startup"
else
export SAGE_ATTENTION_BUILT=0
export SAGE_ATTENTION_AVAILABLE=0
@ -334,6 +330,28 @@ PY
exec runuser -u "${APP_USER}" -- "$0" "$@"
fi
# Early PATH and Python user-site for the runtime user (before any pip work)
export PATH="$HOME/.local/bin:$PATH"
pyver="$(python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
export PYTHONPATH="$HOME/.local/lib/python${pyver}/site-packages:${PYTHONPATH:-}"
# Safer parallelism for extension builds to avoid OOM
detect_jobs() {
local cpu="$(nproc 2>/dev/null || echo 2)"
local mem_mb="$(awk '/MemTotal:/ {print int($2/1024)}' /proc/meminfo 2>/dev/null || echo 0)"
local mem_jobs=2
if [ "$mem_mb" -gt 0 ]; then
if [ "$mem_mb" -ge 32768 ]; then mem_jobs=4; else mem_jobs=2; fi
fi
local cap="$mem_jobs"
if [ "$cpu" -lt "$cap" ]; then cap="$cpu"; fi
[ "$cap" -lt 1 ] && cap=1
echo "$cap"
}
export MAX_JOBS="${SAGEATTENTION_MAX_JOBS:-$(detect_jobs)}"
export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-$MAX_JOBS}"
log "Capping extension build parallelism: MAX_JOBS=$MAX_JOBS, CMAKE_BUILD_PARALLEL_LEVEL=$CMAKE_BUILD_PARALLEL_LEVEL"
# Setup Sage Attention for detected GPU configuration
setup_sage_attention
@ -348,11 +366,6 @@ elif [ ! -d "$CUSTOM_NODES_DIR/ComfyUI-Manager" ]; then
git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager.git "$CUSTOM_NODES_DIR/ComfyUI-Manager" || true
fi
# User-site PATHs for --user installs (custom nodes)
export PATH="$HOME/.local/bin:$PATH"
pyver="$(python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
export PYTHONPATH="$HOME/.local/lib/python${pyver}/site-packages:${PYTHONPATH:-}"
# First-run detection for custom node deps (with override)
RUN_NODE_INSTALL=0
if [ ! -f "$FIRST_RUN_FLAG" ]; then
@ -379,8 +392,6 @@ if [ "$RUN_NODE_INSTALL" = "1" ]; then
done < <(find "$CUSTOM_NODES_DIR" -maxdepth 2 -type f -iname 'pyproject.toml' -not -path '*/ComfyUI-Manager/*' -print0)
python -m pip check || true
# Mark first run complete (or keep flag if already set)
touch "$FIRST_RUN_FLAG" || true
fi