#!/bin/bash
# TP1 v5 sweep — Option D: D-side admission for seed/reseed.
#
# v4 (cap=16) still saw 35% session-cap fallback because the local soft_cap
# evaluates min(16, usable_capacity_tokens / target_tokens) and target_tokens
# (= input + output) is 50-100K in agentic workloads, giving cap = 1-2.
#
# v5 makes worker admission_mode authoritative for ALL admission decisions
# (direct_append AND seed/reseed). Replay calls D's
# /session_cache/admit_direct_append with mode={direct_append|seed} and
# defers to D's KV pool availability + LRU eviction. Replay's local
# _decode_session_soft_cap is bypassed entirely under worker mode.
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-v5-optD
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
}

save_result() {
  local label=$1
  local run_dir=$2
  log "=== $label COMPLETED ==="
  if [ -f "$run_dir/request-metrics.jsonl.summary.json" ]; then
    log "Summary:"
    cat "$run_dir/request-metrics.jsonl.summary.json" >> $RESULTS_FILE
    echo "" >> $RESULTS_FILE
    cp "$run_dir/request-metrics.jsonl.summary.json" "$OUTPUT/${label}_summary.json"
    cp "$run_dir/request-metrics.jsonl" "$OUTPUT/${label}_metrics.jsonl"
    log "Saved to $OUTPUT/${label}_summary.json + ${label}_metrics.jsonl"
  else
    log "WARNING: No summary file found in $run_dir"
  fi
}

log "Starting TP1 v5 sweep (Option D: D-side seed admission)"
log "Model: $MODEL"
log "Trace: $TRACE (4449 requests, 52 sessions)"
log "Key change: worker admission_mode now drives seed/reseed via D's admit endpoint"

########################################
# Experiment 1: 1P + 7D KVC kv-aware Option D
########################################
log ""
log "=== [EXP1] 1P7D KVC kv-aware Option D ==="
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 7 \
  --prefill-tp-size 1 --decode-tp-size 1 \
  --prefill-gpu-ids 0 --decode-gpu-ids 1,2,3,4,5,6,7 \
  --transfer-backend mooncake \
  --gpu-budget 8 \
  --time-scale 10 \
  --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

EXP1_DIR=$(ls -td $OUTPUT/kvcache-centric-*/ 2>/dev/null | head -1)
save_result "exp1_1p7d_kvc_optD" "$EXP1_DIR"

########################################
# Experiment 2: 2P + 6D KVC kv-aware Option D
########################################
log ""
log "=== [EXP2] 2P6D KVC kv-aware Option D ==="
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 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 \
  --time-scale 10 \
  --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

EXP2_DIR=$(ls -td $OUTPUT/kvcache-centric-*/ 2>/dev/null | head -1)
save_result "exp2_2p6d_kvc_optD" "$EXP2_DIR"

log ""
log "=== ALL TP1 V5 SWEEP EXPERIMENTS DONE ==="
