Comfy-Preview-2 / entrypoint.sh
rikunarita's picture
Create entrypoint.sh
82b0c6a verified
#!/bin/bash
# ---------------------------------------------------------------
# 起動時に「このコンテナに割り当てられた」CPU 数を正確に検出する。
#
# nproc --all はホストの全 CPU 数を返すため NG。
# nproc (引数なし) は cpuset 制限は尊重するが、
# HuggingFace のように --cpus=2 (CFS クォータ) で制限している場合は
# やはりホスト値を返してしまう。
#
# 唯一正確な方法は cgroup のクォータファイルを直接読む:
# cgroup v2: /sys/fs/cgroup/cpu.max → "quota period"
# cgroup v1: /sys/fs/cgroup/cpu/cpu.cfs_quota_us
# /sys/fs/cgroup/cpu/cpu.cfs_period_us
# ---------------------------------------------------------------
# cgroup v2 を優先して確認
if [ -f /sys/fs/cgroup/cpu.max ]; then
CPU_MAX=$(cat /sys/fs/cgroup/cpu.max)
QUOTA=$(echo "$CPU_MAX" | awk '{print $1}')
PERIOD=$(echo "$CPU_MAX" | awk '{print $2}')
if [ "$QUOTA" = "max" ]; then
# クォータ無制限 → cpuset ベースの実コア数にフォールバック
NCPU=$(nproc)
else
NCPU=$(( QUOTA / PERIOD ))
fi
# cgroup v1 にフォールバック
elif [ -f /sys/fs/cgroup/cpu/cpu.cfs_quota_us ]; then
QUOTA=$(cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us)
PERIOD=$(cat /sys/fs/cgroup/cpu/cpu.cfs_period_us)
if [ "$QUOTA" -le 0 ]; then
NCPU=$(nproc)
else
NCPU=$(( QUOTA / PERIOD ))
fi
else
# cgroup が読めない環境
NCPU=$(nproc)
fi
# 最低 1 を保証
[ "$NCPU" -lt 1 ] && NCPU=1
echo "[entrypoint] 割り当て CPU 数: ${NCPU}"
# [修正7] jemalloc ランタイム設定を NCPU 確定後に動的セット。
#
# narenas:${NCPU}
# アリーナ数を割り当て vCPU 数 (HF 無料枠 = 2) に合わせる。
# デフォルト (論理コア数 × 4) のままだと、アリーナが独立して
# メモリを管理するため断片化が倍増する。
#
# background_thread:true
# バックグラウンドスレッドが非同期でダーティページを回収する。
# アプリスレッドのレイテンシに影響せず、推論間の RSS スパイクを平滑化。
#
# dirty_decay_ms:2000
# 解放済みダーティページを 2 秒以内に OS へ返却 (デフォルト 10 秒)。
# 大テンソルの解放後に RAM が素早く空くため次の推論がOOMしにくい。
#
# muzzy_decay_ms:0
# decommit 済みページを即座に返却 (デフォルト 10 秒)。
# 仮想メモリのフットプリントを最小化し、RSS の上限を下げる。
export MALLOC_CONF="narenas:${NCPU},background_thread:true,dirty_decay_ms:2000,muzzy_decay_ms:0"
# PyTorch / OpenMP intra-op スレッド数
export OMP_NUM_THREADS=${NCPU}
# Intel MKL スレッド数 (PyTorch が MKL-DNN を使う場合)
export MKL_NUM_THREADS=${NCPU}
# OpenBLAS スレッド数 (numpy 等が内部で使用)
export OPENBLAS_NUM_THREADS=${NCPU}
# numexpr スレッド数 (numpy 高度演算の補助ライブラリ)
export NUMEXPR_NUM_THREADS=${NCPU}
# BLIS スレッド数 (PyTorch のビルドによっては使用)
export BLIS_NUM_THREADS=${NCPU}
exec "$@"