26S30n

#!/usr/bin/env bash
# ==========================================================================
# STAGE 3a-8: Prefix Caching True A/B Test
#
# 목적:
#   동일한 32K prompt를 반복 요청하여 Prefix Caching의 순수 효과 측정
#
# Workload:
#   Total input      = 32,768 tokens
#   Fixed prefix     = 32,000 tokens
#   Random suffix    =    768 tokens
#   Output           =    256 tokens
#   Requests         =    100 / run
#   Concurrency      =      8
#
# 핵심:
#   random-prefix-len은 "전체 input 중 prefix 길이"가 아니라
#   random context 앞에 추가되는 fixed prefix이므로,
#
#       random-input-len = 768
#       random-prefix-len = 32000
#
#   으로 설정하여 총 input = 32000 + 768 = 32768로 맞춘다.
#
#   vLLM metrics:
#     vllm:prefix_cache_queries
#     vllm:prefix_cache_hits
#     vllm:prompt_tokens_cached
#
#   를 함께 수집하여 실제 cache hit 여부를 검증한다.
# ==========================================================================

set -euo pipefail

# --------------------------------------------------------------------------
# Offline environment
# --------------------------------------------------------------------------
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export HF_DATASETS_OFFLINE=1

# --------------------------------------------------------------------------
# Model
# --------------------------------------------------------------------------
LOCAL_MODEL_ROOT="${LOCAL_MODEL_ROOT:-/mnt/local-nvme-cache/models}"
MODEL_70B="${MODEL_70B:-${LOCAL_MODEL_ROOT}/Meta-Llama-3.1-70B-Instruct}"

[ -d "${MODEL_70B}" ] || {
    echo "모델 없음: ${MODEL_70B}"
    exit 1
}

# --------------------------------------------------------------------------
# vLLM
# --------------------------------------------------------------------------
TP_SIZE="${TP_SIZE:-8}"
VLLM_PORT="${VLLM_PORT:-8000}"
SERVER_HOST="${SERVER_HOST:-127.0.0.1}"

# --------------------------------------------------------------------------
# Benchmark workload
# --------------------------------------------------------------------------
TOTAL_INPUT_LEN=32768
PREFIX_LEN=32000
RANDOM_SUFFIX_LEN=$((TOTAL_INPUT_LEN - PREFIX_LEN))
OUTPUT_LEN=256

NUM_PROMPTS="${NUM_PROMPTS:-100}"
MAX_CONCURRENCY="${MAX_CONCURRENCY:-8}"

# 동일한 synthetic dataset 생성
BENCH_SEED="${BENCH_SEED:-12345}"

# 각 조건 독립 반복 횟수
REPEATS="${REPEATS:-2}"

# --------------------------------------------------------------------------
# Logging
# --------------------------------------------------------------------------
LOG_DIR="/var/log/b300_validation/stage3a8_prefixab_$(date +%Y%m%d_%H%M%S)"
mkdir -p "${LOG_DIR}"

log() {
    echo -e "\n\033[1;32m[$(date '+%Y-%m-%d %H:%M:%S')] $1\033[0m"
}

warn() {
    echo -e "\033[1;33mWARNING: $1\033[0m"
}

source ~/vllm-bench-env/bin/activate

# --------------------------------------------------------------------------
# Cleanup
# --------------------------------------------------------------------------
SERVER_PID=""

cleanup_server() {
    if [[ -n "${SERVER_PID}" ]]; then
        kill -TERM "${SERVER_PID}" 2>/dev/null || true
        wait "${SERVER_PID}" 2>/dev/null || true
        SERVER_PID=""
    fi
}

trap cleanup_server EXIT INT TERM

# --------------------------------------------------------------------------
# Wait for server
# --------------------------------------------------------------------------
wait_server() {
    local timeout="${1:-600}"
    local elapsed=0

    log "vLLM server 준비 대기..."

    until curl -sf \
        "http://${SERVER_HOST}:${VLLM_PORT}/v1/models" \
        >/dev/null 2>&1
    do
        if [[ -n "${SERVER_PID}" ]]; then
            kill -0 "${SERVER_PID}" 2>/dev/null || {
                echo "vLLM 서버 기동 실패"
                return 1
            }
        fi

        sleep 5
        elapsed=$((elapsed + 5))

        if (( elapsed >= timeout )); then
            echo "vLLM server readiness timeout (${timeout}s)"
            return 1
        fi
    done

    log "vLLM server READY (${elapsed}s)"
}

# --------------------------------------------------------------------------
# Collect Prometheus metrics
# --------------------------------------------------------------------------
collect_metrics() {
    local output_file="$1"

    curl -sf \
        "http://${SERVER_HOST}:${VLLM_PORT}/metrics" \
        > "${output_file}" || {
            warn "Prometheus metrics 수집 실패: ${output_file}"
            touch "${output_file}"
        }
}

# --------------------------------------------------------------------------
# Extract a counter from Prometheus text
# Handles optional labels.
# --------------------------------------------------------------------------
metric_value() {
    local file="$1"
    local metric="$2"

    awk -v metric="${metric}" '
        $1 ~ ("^" metric "($|\\{)") {
            print $NF
            exit
        }
    ' "${file}" 2>/dev/null || true
}

# --------------------------------------------------------------------------
# Start server
# --------------------------------------------------------------------------
start_server() {
    local tag="$1"
    local cache_flag="$2"
    local server_log="${LOG_DIR}/server_${tag}.log"

    log "============================================================"
    log "서버 기동: ${tag}"
    log "Prefix caching: ${cache_flag}"
    log "============================================================"

    python3 -m vllm.entrypoints.openai.api_server \
        --model "${MODEL_70B}" \
        --served-model-name "llama-3.1-70b" \
        --tensor-parallel-size "${TP_SIZE}" \
        --gpu-memory-utilization 0.90 \
        --max-model-len 65536 \
        --enable-metrics \
        ${cache_flag} \
        --port "${VLLM_PORT}" \
        > "${server_log}" 2>&1 &

    SERVER_PID=$!

    wait_server 900
}

# --------------------------------------------------------------------------
# Warm-up
#
# IMPORTANT:
# random dataset is deterministic for a fixed seed.
#
# Warm-up uses exactly the same seed as the measured run so that the first
# generated prompt is the same fixed-prefix workload.
# --------------------------------------------------------------------------
warmup() {
    local tag="$1"

    log "[${tag}] Prefix cache warm-up"

    vllm bench serve \
        --backend openai-chat \
        --endpoint /v1/chat/completions \
        --host "${SERVER_HOST}" \
        --port "${VLLM_PORT}" \
        --model "llama-3.1-70b" \
        --tokenizer "${MODEL_70B}" \
        --dataset-name random \
        --random-input-len "${RANDOM_SUFFIX_LEN}" \
        --random-output-len "${OUTPUT_LEN}" \
        --random-prefix-len "${PREFIX_LEN}" \
        --random-range-ratio 0 \
        --num-prompts 1 \
        --max-concurrency 1 \
        --request-rate inf \
        --seed "${BENCH_SEED}" \
        --ignore-eos \
        > "${LOG_DIR}/warmup_${tag}.log" 2>&1 || true

    sleep 2
}

# --------------------------------------------------------------------------
# Run benchmark
# --------------------------------------------------------------------------
run_benchmark() {
    local tag="$1"

    local result_file="${LOG_DIR}/${tag}.json"
    local bench_log="${LOG_DIR}/${tag}.log"

    log "[${tag}] Benchmark 시작"
    log "  total input      : ${TOTAL_INPUT_LEN}"
    log "  fixed prefix     : ${PREFIX_LEN}"
    log "  random suffix    : ${RANDOM_SUFFIX_LEN}"
    log "  output           : ${OUTPUT_LEN}"
    log "  requests         : ${NUM_PROMPTS}"
    log "  concurrency      : ${MAX_CONCURRENCY}"

    vllm bench serve \
        --backend openai-chat \
        --endpoint /v1/chat/completions \
        --host "${SERVER_HOST}" \
        --port "${VLLM_PORT}" \
        --model "llama-3.1-70b" \
        --tokenizer "${MODEL_70B}" \
        --dataset-name random \
        --random-input-len "${RANDOM_SUFFIX_LEN}" \
        --random-output-len "${OUTPUT_LEN}" \
        --random-prefix-len "${PREFIX_LEN}" \
        --random-range-ratio 0 \
        --num-prompts "${NUM_PROMPTS}" \
        --max-concurrency "${MAX_CONCURRENCY}" \
        --request-rate inf \
        --seed "${BENCH_SEED}" \
        --ignore-eos \
        --save-result \
        --result-dir "${LOG_DIR}" \
        --result-filename "$(basename "${result_file}")" \
        2>&1 | tee "${bench_log}"

    # vLLM bench may generate a different result path/name depending on
    # installed version. Verify existence.
    if [[ ! -f "${result_file}" ]]; then
        warn "Expected result file not found: ${result_file}"
    fi
}

# --------------------------------------------------------------------------
# Run one complete case
# --------------------------------------------------------------------------
run_case() {
    local condition="$1"
    local cache_flag="$2"
    local repeat="$3"

    local tag="${condition}_r${repeat}"

    start_server "${tag}" "${cache_flag}"

    # Initial metrics
    collect_metrics "${LOG_DIR}/metrics_${tag}_before.txt"

    # Warm-up
    warmup "${tag}"

    # Metrics after warm-up / before measured workload
    collect_metrics "${LOG_DIR}/metrics_${tag}_warmup.txt"

    # Main benchmark
    run_benchmark "${tag}"

    # Final metrics
    collect_metrics "${LOG_DIR}/metrics_${tag}_after.txt"

    # Extract cache metrics
    local q0 q1 h0 h1 c0 c1

    q0=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_warmup.txt" \
        "vllm:prefix_cache_queries")
    q1=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_after.txt" \
        "vllm:prefix_cache_queries")

    h0=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_warmup.txt" \
        "vllm:prefix_cache_hits")
    h1=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_after.txt" \
        "vllm:prefix_cache_hits")

    c0=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_warmup.txt" \
        "vllm:prompt_tokens_cached")
    c1=$(metric_value \
        "${LOG_DIR}/metrics_${tag}_after.txt" \
        "vllm:prompt_tokens_cached")

    python3 - "${tag}" \
        "${q0:-0}" "${q1:-0}" \
        "${h0:-0}" "${h1:-0}" \
        "${c0:-0}" "${c1:-0}" \
        > "${LOG_DIR}/metrics_${tag}.json" <<'PYEOF'
import json
import sys

tag = sys.argv[1]

q0, q1 = float(sys.argv[2]), float(sys.argv[3])
h0, h1 = float(sys.argv[4]), float(sys.argv[5])
c0, c1 = float(sys.argv[6]), float(sys.argv[7])

queries = max(0, q1 - q0)
hits = max(0, h1 - h0)
cached_tokens = max(0, c1 - c0)

hit_rate = hits / queries if queries > 0 else 0.0

print(json.dumps({
    "tag": tag,
    "prefix_cache_queries": queries,
    "prefix_cache_hits": hits,
    "prefix_cache_hit_rate": hit_rate,
    "prompt_tokens_cached": cached_tokens,
}, indent=2))
PYEOF

    log "[${tag}] Metrics:"
    cat "${LOG_DIR}/metrics_${tag}.json"

    cleanup_server
    sleep 5
}

# ==========================================================================
# Main
# ==========================================================================

log "============================================================"
log "STAGE 3a-8 Prefix Caching A/B Test"
log "============================================================"

log "Workload:"
log "  Total input tokens : ${TOTAL_INPUT_LEN}"
log "  Fixed prefix       : ${PREFIX_LEN}"
log "  Random suffix      : ${RANDOM_SUFFIX_LEN}"
log "  Output tokens      : ${OUTPUT_LEN}"
log "  Requests           : ${NUM_PROMPTS}"
log "  Concurrency        : ${MAX_CONCURRENCY}"
log "  Seed               : ${BENCH_SEED}"
log "  Repeats            : ${REPEATS}"

# Sanity check
if (( PREFIX_LEN + RANDOM_SUFFIX_LEN != TOTAL_INPUT_LEN )); then
    echo "ERROR: input length mismatch"
    exit 1
fi

# --------------------------------------------------------------------------
# Run A/B
# --------------------------------------------------------------------------

for repeat in $(seq 1 "${REPEATS}"); do
    run_case "prefix_cache_on" "--enable-prefix-caching" "${repeat}"
    run_case "prefix_cache_off" "--no-enable-prefix-caching" "${repeat}"
done

# ==========================================================================
# Result analysis
# ==========================================================================

log "=== 결과 분석 ==="

python3 - "${LOG_DIR}" <<'PYEOF'

import glob
import json
import os
import statistics
import re
import sys

log_dir = sys.argv[1]

def load_json(path):
    try:
        with open(path) as f:
            return json.load(f)
    except Exception:
        return {}

def find_result(tag):
    candidates = [
        os.path.join(log_dir, f"{tag}.json"),
    ]

    for path in candidates:
        if os.path.exists(path):
            return load_json(path)

    return {}

def find_metric(tag):
    path = os.path.join(log_dir, f"metrics_{tag}.json")
    return load_json(path)

def get_num(data, *keys):
    for key in keys:
        value = data.get(key)
        if isinstance(value, (int, float)):
            return float(value)
    return None

# ----------------------------------------------------------------------
# Extract benchmark results
# ----------------------------------------------------------------------

on_results = []
off_results = []

for path in sorted(glob.glob(os.path.join(log_dir, "prefix_cache_on_r*.json"))):
    data = load_json(path)
    if data:
        on_results.append(data)

for path in sorted(glob.glob(os.path.join(log_dir, "prefix_cache_off_r*.json"))):
    data = load_json(path)
    if data:
        off_results.append(data)

def summarize(results):
    fields = [
        "median_ttft_ms",
        "p95_ttft_ms",
        "p99_ttft_ms",
        "mean_ttft_ms",
        "median_tpot_ms",
        "p95_tpot_ms",
        "request_throughput",
        "output_throughput",
        "total_token_throughput",
    ]

    out = {}

    for field in fields:
        values = [
            get_num(x, field)
            for x in results
            if get_num(x, field) is not None
        ]

        if values:
            out[field] = statistics.mean(values)

    return out

on = summarize(on_results)
off = summarize(off_results)

# ----------------------------------------------------------------------
# Prefix cache metrics
# ----------------------------------------------------------------------

on_cache = []

for path in sorted(glob.glob(
    os.path.join(log_dir, "metrics_prefix_cache_on_r*.json")
)):
    data = load_json(path)
    if data:
        on_cache.append(data)

off_cache = []

for path in sorted(glob.glob(
    os.path.join(log_dir, "metrics_prefix_cache_off_r*.json")
)):
    data = load_json(path)
    if data:
        off_cache.append(data)

def avg_metric(items, key):
    values = [
        float(x[key])
        for x in items
        if key in x
    ]
    return statistics.mean(values) if values else 0.0

on_hit_rate = avg_metric(on_cache, "prefix_cache_hit_rate")
on_cached_tokens = avg_metric(on_cache, "prompt_tokens_cached")

off_hit_rate = avg_metric(off_cache, "prefix_cache_hit_rate")
off_cached_tokens = avg_metric(off_cache, "prompt_tokens_cached")

# ----------------------------------------------------------------------
# Speedup
# ----------------------------------------------------------------------

off_ttft = off.get("median_ttft_ms", 0)
on_ttft = on.get("median_ttft_ms", 0)

speedup = (
    off_ttft / on_ttft
    if on_ttft and on_ttft > 0
    else 0
)

improvement_pct = (
    (1 - on_ttft / off_ttft) * 100
    if off_ttft and off_ttft > 0
    else 0
)

# ----------------------------------------------------------------------
# Interpretation
# ----------------------------------------------------------------------

if speedup >= 8:
    interpretation = "매우 큰 Prefix Caching 효과"
elif speedup >= 4:
    interpretation = "큰 Prefix Caching 효과"
elif speedup >= 2:
    interpretation = "명확한 Prefix Caching 효과"
elif speedup >= 1.5:
    interpretation = "효과가 있으나 상대적으로 제한적"
else:
    interpretation = "효과가 작음 — cache hit 및 workload 조건 재확인 필요"

# ----------------------------------------------------------------------
# Markdown report
# ----------------------------------------------------------------------

report = f"""# Stage 3a-8 Prefix Caching A/B 결과

## Workload

| 항목 | 값 |
|---|---:|
| Total input | 32,768 tokens |
| Fixed prefix | 32,000 tokens |
| Random suffix | 768 tokens |
| Output | 256 tokens |
| Requests / run | {os.environ.get("NUM_PROMPTS", "100")} |
| Concurrency | {os.environ.get("MAX_CONCURRENCY", "8")} |
| Repeats | {os.environ.get("REPEATS", "2")} |
| Seed | {os.environ.get("BENCH_SEED", "12345")} |

> Prefix = 32,000 tokens + random suffix = 768 tokens,
> therefore total input = 32,768 tokens.

## TTFT

| Metric | Cache OFF | Cache ON | 결과 |
|---|---:|---:|---:|
| Median TTFT | {off.get("median_ttft_ms", 0):.1f} ms | {on.get("median_ttft_ms", 0):.1f} ms | {speedup:.2f}x |
| P95 TTFT | {off.get("p95_ttft_ms", 0):.1f} ms | {on.get("p95_ttft_ms", 0):.1f} ms | - |
| P99 TTFT | {off.get("p99_ttft_ms", 0):.1f} ms | {on.get("p99_ttft_ms", 0):.1f} ms | - |

### Median TTFT improvement

**{improvement_pct:.1f}%**

**Speedup: {speedup:.2f}x**

## Prefix Cache Metrics

| Metric | Cache OFF | Cache ON |
|---|---:|---:|
| Prefix cache hit rate | {off_hit_rate:.2%} | {on_hit_rate:.2%} |
| Prompt tokens cached | {off_cached_tokens:,.0f} | {on_cached_tokens:,.0f} |

## Other Performance Metrics

| Metric | Cache OFF | Cache ON |
|---|---:|---:|
| Median TPOT | {off.get("median_tpot_ms", 0):.2f} ms | {on.get("median_tpot_ms", 0):.2f} ms |
| Request throughput | {off.get("request_throughput", 0):.2f} req/s | {on.get("request_throughput", 0):.2f} req/s |
| Output throughput | {off.get("output_throughput", 0):.2f} tok/s | {on.get("output_throughput", 0):.2f} tok/s |
| Total token throughput | {off.get("total_token_throughput", 0):.2f} tok/s | {on.get("total_token_throughput", 0):.2f} tok/s |

## Interpretation

**{interpretation}**

Reference interpretation for this POC:

- `< 1.5x` : 효과가 작음
- `1.5x ~ 2x` : 제한적인 효과
- `2x ~ 4x` : 명확한 효과
- `4x ~ 8x` : 큰 효과
- `>= 8x` : 매우 큰 효과

These are **POC interpretation thresholds**, not official vLLM pass/fail criteria.

## Validation Criteria

### Functional

- [ ] 동일한 32K fixed prefix가 모든 request에서 사용됨
- [ ] Cache ON에서 prefix cache hit가 관찰됨
- [ ] Cache OFF에서는 prefix cache hit가 발생하지 않음
- [ ] Cache ON/OFF의 output length가 동일함
- [ ] Request failure가 없어야 함

### Performance

- [ ] Median TTFT 감소
- [ ] P95/P99 TTFT도 악화되지 않는지 확인
- [ ] Prefix cache hit와 TTFT 개선이 함께 관찰되는지 확인
- [ ] TPOT/ITL이 불필요하게 악화되지 않는지 확인

## Important Note

이 테스트는 **32K fixed-prefix 반복이라는 Prefix Caching에 매우 유리한 synthetic workload**이다.

따라서 이 결과를 실제 서비스의 예상 성능 향상률로 직접 해석해서는 안 된다.

실제 workload 효과는 다음 단계에서:

- Prefix length: 8K / 16K / 32K
- Prefix reuse ratio: 25% / 50% / 75% / 100%
- Concurrency: 1 / 2 / 4 / 8 / 16

등으로 추가 검증하는 것이 적절하다.
"""

report_path = os.path.join(
    log_dir,
    "prefix_cache_ab_report.md"
)

with open(report_path, "w") as f:
    f.write(report)

print(report)
print()
print(f"Report: {report_path}")

PYEOF

log "============================================================"
log "Prefix Caching A/B 완료"
log "Report: ${LOG_DIR}/prefix_cache_ab_report.md"
log "============================================================"