#!/usr/bin/env bash
# Real Ali workload sweep for KVC pd-hybrid.
#
# This script expects a prebuilt sample trace and replays it exactly for every
# mechanism.  It intentionally keeps pool polling disabled for performance runs.
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$REPO_ROOT"

MODEL=${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}
TRACE=${TRACE:-outputs/real-ali-kvc-iter/samples-balanced/ali-kvc-fit-smallappend.jsonl}
OUT_ROOT=${OUT_ROOT:-outputs/real-ali-kvc-iter/runs}
TIME_SCALE=${TIME_SCALE:-1}
CONCURRENCY=${CONCURRENCY:-32}
REQUEST_TIMEOUT_S=${REQUEST_TIMEOUT_S:-300}
STACK_TIMEOUT_S=${STACK_TIMEOUT_S:-1200}
RUNS=${RUNS:-"dp kvc_bp"}
EXTRA_SERVER_ARGS=${EXTRA_SERVER_ARGS:-}
PREFILL_EXTRA_SERVER_ARGS=${PREFILL_EXTRA_SERVER_ARGS:-}
DECODE_EXTRA_SERVER_ARGS=${DECODE_EXTRA_SERVER_ARGS:-}
KVC_SEED_MIN_TURN_ID=${KVC_SEED_MIN_TURN_ID:-1}
KVC_SEED_ONLY_MULTITURN=${KVC_SEED_ONLY_MULTITURN:-0}

mkdir -p "$OUT_ROOT"
LOG="$OUT_ROOT/sweep.log"

log() {
  echo "[$(date '+%F %T')] $*" | tee -a "$LOG"
}

common_args=(
  --trace "$TRACE"
  --model-path "$MODEL"
  --output-root "$OUT_ROOT"
  --use-trace-as-sample
  --time-scale "$TIME_SCALE"
  --concurrency-limit "$CONCURRENCY"
  --timeout-s "$STACK_TIMEOUT_S"
  --request-timeout-s "$REQUEST_TIMEOUT_S"
)
if [[ -n "$EXTRA_SERVER_ARGS" ]]; then
  common_args+=(--extra-server-args "$EXTRA_SERVER_ARGS")
fi
if [[ -n "$PREFILL_EXTRA_SERVER_ARGS" ]]; then
  common_args+=(--prefill-extra-server-args "$PREFILL_EXTRA_SERVER_ARGS")
fi
if [[ -n "$DECODE_EXTRA_SERVER_ARGS" ]]; then
  common_args+=(--decode-extra-server-args "$DECODE_EXTRA_SERVER_ARGS")
fi

kvc_args=(
  "${common_args[@]}"
  --mechanism kvcache-centric
  --policy kv-aware
  --prefill-workers 2
  --decode-workers 6
  --prefill-tp-size 1
  --decode-tp-size 1
  --prefill-gpu-ids 0,1
  --decode-gpu-ids 2,3,4,5,6,7
  --transfer-backend mooncake
  --gpu-budget 8
  --kvcache-admission-mode worker
  --kvcache-seed-min-turn-id "$KVC_SEED_MIN_TURN_ID"
  --kvcache-seed-max-inflight-decode -1
  --kvcache-prefill-backup-policy release-after-transfer
  --kvcache-prefill-priority-eviction
)
if [[ "$KVC_SEED_ONLY_MULTITURN" == "1" ]]; then
  kvc_args+=(--kvcache-seed-only-multiturn-sessions)
fi

run_dp() {
  log "=== DP cache-aware baseline: 8 direct workers ==="
  uv run agentic-pd-hybrid benchmark-live \
    "${common_args[@]}" \
    --mechanism pd-colo \
    --policy kv-aware \
    --prefill-workers 0 \
    --decode-workers 0 \
    --direct-workers 8 \
    --direct-tp-size 1 \
    --direct-gpu-ids 0,1,2,3,4,5,6,7 \
    --gpu-budget 8
}

run_pd_disagg() {
  log "=== PD-disaggregation baseline: 2P6D ==="
  uv run agentic-pd-hybrid benchmark-live \
    "${common_args[@]}" \
    --mechanism pd-disaggregation \
    --policy kv-aware \
    --prefill-workers 2 \
    --decode-workers 6 \
    --prefill-tp-size 1 \
    --decode-tp-size 1 \
    --prefill-gpu-ids 0,1 \
    --decode-gpu-ids 2,3,4,5,6,7 \
    --transfer-backend mooncake \
    --gpu-budget 8
}

run_pd_sticky() {
  log "=== PD-disaggregation sticky baseline: 2P6D ==="
  uv run agentic-pd-hybrid benchmark-live \
    "${common_args[@]}" \
    --mechanism pd-disaggregation \
    --policy sticky \
    --prefill-workers 2 \
    --decode-workers 6 \
    --prefill-tp-size 1 \
    --decode-tp-size 1 \
    --prefill-gpu-ids 0,1 \
    --decode-gpu-ids 2,3,4,5,6,7 \
    --transfer-backend mooncake \
    --gpu-budget 8
}

run_kvc() {
  log "=== KVC baseline: 2P6D worker admission, no backpressure ==="
  uv run agentic-pd-hybrid benchmark-live "${kvc_args[@]}"
}

run_kvc_bp() {
  log "=== KVC candidate: 2P6D worker admission + backpressure ==="
  uv run agentic-pd-hybrid benchmark-live \
    "${kvc_args[@]}" \
    --enable-backpressure \
    --backpressure-max-pause-s 2.0
}

summarize_latest() {
  log "=== Latest summaries ==="
  find "$OUT_ROOT" -maxdepth 2 -name 'request-metrics.jsonl.summary.json' -print \
    | sort \
    | while read -r summary; do
        python - "$summary" <<'PY'
import json, sys
p=sys.argv[1]
d=json.load(open(p))
lat=d.get("latency_stats_s") or {}
tt=d.get("ttft_stats_s") or {}
em=d.get("execution_modes") or {}
print(p)
print("  reqs", d.get("request_count"), "errors", d.get("error_count"), "trunc", d.get("truncated_request_count"))
print("  lat mean/p50/p90/p99", lat.get("mean"), lat.get("p50"), lat.get("p90"), lat.get("p99"))
print("  ttft mean/p50/p90", tt.get("mean"), tt.get("p50"), tt.get("p90"))
print("  modes", em)
PY
      done | tee -a "$LOG"
}

log "Trace: $TRACE"
log "Model: $MODEL"
log "Runs: $RUNS | time-scale=$TIME_SCALE concurrency=$CONCURRENCY | kvc-seed-min-turn-id=$KVC_SEED_MIN_TURN_ID | kvc-seed-only-multiturn=$KVC_SEED_ONLY_MULTITURN"

for run in $RUNS; do
  case "$run" in
    dp) run_dp ;;
    pd) run_pd_disagg ;;
    pd_sticky) run_pd_sticky ;;
    kvc) run_kvc ;;
    kvc_bp) run_kvc_bp ;;
    *) log "Unknown run name: $run"; exit 2 ;;
  esac
done

summarize_latest
log "DONE"