#!/bin/bash

cd /data/bts-ai-compose/api/binzib_chatbot_v2

source /home/was/project/.venv/bin/activate

COMPOSE_FILE="docker-compose.infra.yml"
RESULTS_DIR="test/results"
SUMMARY_FILE="$RESULTS_DIR/summary.txt"
mkdir -p $RESULTS_DIR

echo #"==================================================================" | tee $SUMMARY_FILE
echo "GPU Memory Utilization 부하 테스트" | tee -a $SUMMARY_FILE
echo #"==================================================================" | tee -a $SUMMARY_FILE

printf "| %-6s | %-10s | %-10s | %-10s | %-10s | %-8s |\n" \
    "GPU%(ms)" "평균", "중간", "95%", "최대", "실패율" |  tee -a $SUMMARY_FILE
echo "|--------|--------|--------|--------|--------|--------|" |  tee -a $SUMMARY_FILE

for UTIL in 0.5 0.6 0.7 0.8 0.85 0.9 0.95; do 
    echo ""
    echo "테스트 시작: gpu-memory-utilization = $UTIL"

    sed -i "s/--gpu-memory-utilization [0-9.]\+/--gpu-memory-utilization $UTIL/" $COMPOSE_FILE

    docker compose -f $COMPOSE_FILE up -d vllm

    echo "vLLM 로드 중..."
    sleep 60 

    for i in {1..10}; do
        if curl -s http://localhost:8005/health > /dev/null 2>&1; then 
            echo "vLLM 세팅 완료"
            break

        fi
        echo "대기중... ($i/10)"
        sleep 5
    done 

    locust -f test/locustfile.py --host http://localhost:8000 --headless -u 10 -r 2 -t 60s --csv=$RESULTS_DIR/gpu_$UTIL --only-summary 2>&1 | tee $RESULTS_DIR/log_$UTIL.txt

    STATS_FILE="$RESULTS_DIR/gpu_${UTIL}_stats.csv"

    if [ -f "$STATS_FILE"]; then

        LINE=$(grep "Aggregated" $STATS_FILE)
        if [ -n "$LINE"]; then
            AVG=$(echo $LINE | cut -d',' -f6)
            MED=$(echo $LINE | cut -d',' -f8)
            P95=$(echo $LINE | cut -d',' -f12)
            MAX=$(echo $LINE | cut -d',' -f7)
            FAIL=$(echo $LINE | cut -d',' -f4)
            TOTAL=$(echo $LINE | cut -d',' -f3)

            if [ "$TOTAL" -gt 0 ] 2>/dev/null; then
                FAIL_RATE=$(echo "scale=2; $FAIL * 100 / $TOTAL" | bc)%
            else
                FAIL_RATE="N/A"
            fi

            printf "| %-6s | %-10s | %-10s | %-10s | %-10s | %-8s |\n" \
                "$UTIL" "$AVG" "$MED" "$P95" "$MAX" "$FAIL_RATE" | tee -a $SUMMARY_FILE

        fi
    fi

    echo "완료: $UTIL"
    sleep 5

done
            
