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