#!/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 "$@"