#!/bin/bash
# Time-scale=1 validation sweep, downscaled to 4 GPUs:
#   - KVC v5 1P3D × N=3   (new data, validates §1/§2 structural claims at real timing)
#   - 4-way DP cache-aware × 1 (sanity baseline at same scale + ts=1)
#
# Goal: per docs/AGENTIC_FIT_ANALYSIS_ZH.md §7 / TEAM_REPORT §2.6 — all v3-v6 KVC
# data was at time-scale=10 (inter-turn gap p50 = 0.25s, vs real 2.5s). This run
# tests whether the gap structurally reverses any conclusion.
#
# CONFIG NOTE: Original experiments used 8 GPUs (2P6D / 8-way DP). This host has
# only 4 H100s available, so we downscale proportionally to 1P3D / 4-way DP.
# Cross-compare against existing 2P6D ts=10 data is confounded by *both*
# time-scale and capacity. Internal comparison (1P3D KVC vs 4DP) at ts=1 is the
# clean signal. §5 (P-side imbalance) is NOT testable here — only 1 P.
#
# Capacity ratio: 3D × ~92K tok = 276K KV pool vs 52 sessions × ~50K peak input
#   working set ≈ 1.5M → ~5.4× overload (vs 2.7× in original 2P6D).
#   Pressure is HIGHER than original; partly offset by ts=1 letting D drain between turns.
#
# Output:
#   outputs/qwen3-30b-tp1-ts1-validation/
#     ├── kvc_1p3d_run{1,2,3}_summary.json
#     ├── kvc_1p3d_run{1,2,3}_metrics.jsonl
#     ├── dp4_summary.json
#     ├── dp4_metrics.jsonl
#     └── kvcache-centric-... / pd-colo-kv-aware-...    (raw run dirs)
#
# Estimated GPU time: KVC ts=1 ≈ 100-180 min/run × 3 = 5-9h
#                     DP ts=1  ≈ 100-120 min × 1     = ~2h
#                     Total                           = 7-11h
set -euo pipefail
cd "$(dirname "$0")/.."

MODEL=/mnt/kzlin/workflow/pd-hybrid/simm-swe-bench/models/Qwen3-30B-A3B-Instruct-2507
TRACE=outputs/qwen35-swebench-50sess.jsonl
OUTPUT=outputs/qwen3-30b-tp1-ts1-validation
VENV_PYTHON=.venv/bin/python
RESULTS_FILE=$OUTPUT/sweep_results.txt

mkdir -p $OUTPUT

log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a $RESULTS_FILE
}

run_kvc_1p3d() {
  local run_idx=$1
  local label="kvc_1p3d_run${run_idx}"
  log ""
  log "=== [KVC ${run_idx}/3] 1P3D KVC kv-aware Option D, time-scale=1 ==="
  PYTHONPATH=src:third_party/sglang/python \
  $VENV_PYTHON -m agentic_pd_hybrid.cli benchmark-live \
    --trace $TRACE \
    --output-root $OUTPUT \
    --mechanism kvcache-centric \
    --policy kv-aware \
    --model-path $MODEL \
    --prefill-workers 1 --decode-workers 3 \
    --prefill-tp-size 1 --decode-tp-size 1 \
    --prefill-gpu-ids 0 --decode-gpu-ids 1,2,3 \
    --transfer-backend mooncake \
    --gpu-budget 4 \
    --time-scale 1 \
    --session-sample-rate 1.0 \
    --target-duration-s 100000 \
    --concurrency-limit 32 \
    --timeout-s 900 \
    --request-timeout-s 300 \
    --kvcache-admission-mode worker \
    --kvcache-seed-min-turn-id 1 \
    --kvcache-seed-max-inflight-decode -1 \
    --kvcache-prefill-backup-policy release-after-transfer \
    --kvcache-prefill-priority-eviction

  local run_dir=$(ls -td $OUTPUT/kvcache-centric-*/ 2>/dev/null | head -1)
  log "=== [KVC ${run_idx}/3] $label COMPLETED ==="
  if [ -f "$run_dir/request-metrics.jsonl.summary.json" ]; then
    cp "$run_dir/request-metrics.jsonl.summary.json" "$OUTPUT/${label}_summary.json"
    cp "$run_dir/request-metrics.jsonl" "$OUTPUT/${label}_metrics.jsonl"
    local errs=$($VENV_PYTHON -c "import json; d=json.load(open('$OUTPUT/${label}_summary.json')); print(d.get('error_count',0))")
    log "  errors = $errs"
    cat "$run_dir/request-metrics.jsonl.summary.json" >> $RESULTS_FILE
    echo "" >> $RESULTS_FILE
  else
    log "WARNING: no summary file in $run_dir"
  fi
}

run_dp4_sanity() {
  local label="dp4"
  log ""
  log "=== [DP] 4-way DP cache-aware sanity, time-scale=1 ==="
  PYTHONPATH=src:third_party/sglang/python \
  $VENV_PYTHON -m agentic_pd_hybrid.cli benchmark-live \
    --trace $TRACE \
    --output-root $OUTPUT \
    --mechanism pd-colo \
    --policy kv-aware \
    --model-path $MODEL \
    --prefill-workers 0 --decode-workers 0 \
    --direct-workers 4 --direct-tp-size 1 \
    --direct-gpu-ids 0,1,2,3 \
    --gpu-budget 4 \
    --time-scale 1 \
    --session-sample-rate 1.0 \
    --target-duration-s 100000 \
    --concurrency-limit 32 \
    --timeout-s 900 \
    --request-timeout-s 300

  local run_dir=$(ls -td $OUTPUT/pd-colo-kv-aware-*/ 2>/dev/null | head -1)
  log "=== [DP] $label COMPLETED ==="
  if [ -f "$run_dir/request-metrics.jsonl.summary.json" ]; then
    cp "$run_dir/request-metrics.jsonl.summary.json" "$OUTPUT/${label}_summary.json"
    cp "$run_dir/request-metrics.jsonl" "$OUTPUT/${label}_metrics.jsonl"
    local errs=$($VENV_PYTHON -c "import json; d=json.load(open('$OUTPUT/${label}_summary.json')); print(d.get('error_count',0))")
    log "  errors = $errs"
    cat "$run_dir/request-metrics.jsonl.summary.json" >> $RESULTS_FILE
    echo "" >> $RESULTS_FILE
  else
    log "WARNING: no summary file in $run_dir"
  fi
}

log "=== TS=1 VALIDATION (4-GPU): KVC 1P3D × N=3 + 4DP × 1 ==="
log "Model: $MODEL"
log "Trace: $TRACE (4449 requests, 52 sessions)"
log "Goal: validate whether ts=10 was the main distortion in v3-v6 KVC vs DP"

# KVC × 3 first (the new data we need); DP last (cheaper sanity at end)
for i in 1 2 3; do
  run_kvc_1p3d $i
done

run_dp4_sanity

log ""
log "=== TS=1 SUMMARY ==="
for label in kvc_1p3d_run1 kvc_1p3d_run2 kvc_1p3d_run3 dp4; do
  if [ -f "$OUTPUT/${label}_summary.json" ]; then
    e=$($VENV_PYTHON -c "import json; d=json.load(open('$OUTPUT/${label}_summary.json')); print(d.get('error_count',0))")
    p50=$($VENV_PYTHON -c "import json; d=json.load(open('$OUTPUT/${label}_summary.json')); print(d.get('latency_stats_s',{}).get('p50','n/a'))")
    log "  ${label}: errors=$e  lat_p50=${p50}s"
  fi
done
log "=== TS=1 ALL DONE ==="