Anima-Preview-3 / entrypoint.sh
CJHauser's picture
Update entrypoint.sh
62e74b1 verified
Raw
History Blame Contribute Delete
1.71 kB
#!/bin/bash
# ---------------------------------------------------------------
# Detect the actual allocated CPU count from cgroup quotas.
# nproc alone returns host CPU count, not the container's share.
# HF free tier uses CFS quota (--cpus=2), so we read cgroup directly.
# ---------------------------------------------------------------
# Try cgroup v2 first
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
NCPU=$(nproc)
else
NCPU=$(( QUOTA / PERIOD ))
fi
# Fallback to 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
NCPU=$(nproc)
fi
# Guarantee at least 1
[ "$NCPU" -lt 1 ] && NCPU=1
echo "[entrypoint] Allocated CPUs: ${NCPU}"
# jemalloc runtime config — set after NCPU is known
# narenas: match arena count to vCPU count (avoids excess fragmentation)
# background_thread: async dirty page reclaim (no latency hit on inference)
# dirty_decay_ms: return freed pages to OS within 2s (default 10s)
# muzzy_decay_ms: immediately decommit pages (keeps RSS low)
export MALLOC_CONF="narenas:${NCPU},background_thread:true,dirty_decay_ms:2000,muzzy_decay_ms:0"
# Thread counts for all major BLAS/compute libs
export OMP_NUM_THREADS=${NCPU}
export MKL_NUM_THREADS=${NCPU}
export OPENBLAS_NUM_THREADS=${NCPU}
export NUMEXPR_NUM_THREADS=${NCPU}
export BLIS_NUM_THREADS=${NCPU}
exec "$@"