# TP×DP matrix experiment for DeepSeek-V4-Flash-INT8 on Kunlun P800 (8 XPUs)
# using SGLang-XPU with the OFFICIAL vendor-recommended launch parameters
# (extracted from the image's /workspace/ds_v4/w8a8/run_server.sh):
#   --ep-size <TP*DP> (expert parallel; probes whether EP makes TP2/DP4 viable)
#   EAGLE speculative decoding (steps=3, topk=1, draft=4)
#   --chunked-prefill-size 8192 --max-prefill-tokens 16384
# Tests SGLang with three parallel configurations:
#   TP=2, DP=4  -> 2 XPUs per replica, 4 replicas
#   TP=4, DP=2  -> 4 XPUs per replica, 2 replicas
#   TP=8, DP=1  -> 8 XPUs, no data parallelism (production config)

EXPERIMENT="dsv4_p800_sglang_tp_dp_official"
MODEL_NAME="DeepSeek-V4-Flash-INT8"
MODEL_PATH="${MODEL_PATH:-/data1/models/DeepSeek-V4-Flash-INT8}"
SERVED_MODEL_NAME="deepseek-v4-flash-int8"

SGLANG_PORT="${SGLANG_PORT:-30014}"

# Dedicated container name so this experiment never touches the
# `sglang-dsv4-flash` container used by the other P800 experiments.
CONTAINER_NAME="sglang-dsv4-flash-tpdp-official"

# Python interpreter inside the P800 SGLang container (used to run the
# benchmark client via `docker exec`).
CONTAINER_PYTHON="${CONTAINER_PYTHON:-/root/miniconda/envs/python310_torch25_cuda/bin/python}"

export XPU_VISIBLE_DEVICES="${XPU_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"

# Runtime working directory for logs, pid files, and tmp. Defaults to a local
# directory under this experiment so the benchmark is self-contained.
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"

# Parallel configurations to test. Format: "TP DP"
# TP2/DP4 is excluded: --ep-size 8 crashes this vendor image when DP>1
# (ZeroDivisionError in data_parallel_controller.py), and --ep-size 2 offers
# no memory relief (~264/2=132 GiB > 96 GiB per XPU), so TP=2 remains
# infeasible for this INT8 model on P800 regardless of EP.
# --ep-size is only passed when DP=1 (official production config TP8/EP8).
# Override with a space-separated list of comma pairs, e.g.
# PARALLEL_CONFIGS_STR="8,1" to run only TP=8/DP=1.
if [[ -n "${PARALLEL_CONFIGS_STR:-}" ]]; then
  declare -a PARALLEL_CONFIGS=()
  for pair in $PARALLEL_CONFIGS_STR; do
    PARALLEL_CONFIGS+=("${pair//,/ }")
  done
else
  declare -a PARALLEL_CONFIGS=(
    "4 2"
    "8 1"
  )
fi

# SGLang server settings, verified for P800 / DeepSeek-V4-Flash-INT8.
CONTEXT_LENGTH="${CONTEXT_LENGTH:-65536}"
MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.8}"

# Official w8a8 parameters (from /workspace/ds_v4/w8a8/run_server.sh).
# EP_SIZE defaults to TP*DP in start_sglang_docker.sh (8 on this machine);
# override here only for single-config probes.
EP_SIZE="${EP_SIZE:-}"
CHUNKED_PREFILL_SIZE="${CHUNKED_PREFILL_SIZE:-8192}"
MAX_PREFILL_TOKENS="${MAX_PREFILL_TOKENS:-16384}"
# EAGLE speculative decoding (official defaults: steps=3, topk=1, draft=4).
SPEC_NUM_STEPS="${SPEC_NUM_STEPS:-3}"
SPEC_EAGLE_TOPK="${SPEC_EAGLE_TOPK:-1}"
SPEC_NUM_DRAFT_TOKENS="${SPEC_NUM_DRAFT_TOKENS:-4}"

DOCKER_IMAGE="${DOCKER_IMAGE:-iregistry.baidu-int.com/xpu/sglang-p800-pd-disagg-0510:20260511_4202}"

# ShareGPT seed dataset for the random workload. DATASET_PATH is the host
# copy (also used by the adaptive search for capacity checks); it is mounted
# into the container at CONTAINER_DATASET_PATH, which is what the benchmark
# client references. Default matches the H20 experiments; the file is
# gitignored, download it into /data1/yy/sskj/datasets/ first.
DATASET_PATH="${DATASET_PATH:-/data1/yy/sskj/datasets/ShareGPT_V3_unfiltered_cleaned_split.json}"
CONTAINER_DATASET_PATH="${CONTAINER_DATASET_PATH:-/workspace/ShareGPT_V3_unfiltered_cleaned_split.json}"

# Matrix and concurrency rules are defined in matrix.json by default.
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR:-.}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"

# Sampling density for concurrency (kept for parity with the H20 experiment;
# generate_scenarios.py currently samples the low/high endpoints only).
export CONCURRENCY_SAMPLES="${CONCURRENCY_SAMPLES:-2}"

# Per-scenario timeout to avoid hangs (seconds).
SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}"

# XPU memory sampling interval (seconds).
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"

# Dry-run mode: if 1, only log the server args and scenario plan without
# starting any server or sending requests.
DRY_RUN="${DRY_RUN:-0}"

# Per-config scenario limit for quick smoke tests. 0 = run all generated
# scenarios.
GRID_LIMIT="${GRID_LIMIT:-0}"