_config.sh

#!/bin/bash
# Shared configuration and helpers for the RankGuide scripts.
#
# Sourced by every other script in this directory. Nothing here runs a job; it
# only sets defaults and defines functions.
#
# Every setting below is overridable from the environment, e.g.
#   CUDA_VISIBLE_DEVICES=0,1 TP_SIZE=2 bash scripts/sweep_grid_steer.sh
# Assigning with `: "${VAR:=default}"` means an exported VAR always wins.

# All relative paths (rankguide.py, vector_500_500/, rankguide/results/) are
# resolved from the TensorRouter/ directory, so run from anywhere.
REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SCRIPT_DIR="${REPO_DIR}/scripts"
cd "$REPO_DIR" || exit 1

# ---------------------------------------------------------------- environment
: "${DATA_DIR:=.}"                                  # root for rankguide/{results,logs}
DATA_DIR="$(realpath -m "$DATA_DIR")"               # absolute, so paths work from any cwd
: "${HF_HOME:=${HOME}/.cache/huggingface}"          # HuggingFace cache
: "${CUDA_VISIBLE_DEVICES:=0,1,2,3}"
export DATA_DIR HF_HOME CUDA_VISIBLE_DEVICES
export VLLM_ALLOW_INSECURE_SERIALIZATION=1

# --------------------------------------------------------------- model serving
: "${BASE_MODEL_NAME:=deepseek-ai/DeepSeek-R1-Distill-Qwen-32B}"   # LRM
: "${SMALL_MODEL_NAME:=deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B}" # SRM
: "${BASE_MODEL_ABBRV:=deepseek-32B}"               # short names, used in result paths
: "${SMALL_MODEL_ABBRV:=deepseek-1.5B}"
: "${BASE_MODEL_PORT:=50001}"
: "${SMALL_MODEL_PORT:=50002}"
: "${TP_SIZE:=2}"                                   # tensor parallel, both models
                                                    # (must be <= number of visible GPUs)
: "${MAX_MODEL_LEN:=20000}"
: "${MAX_NUM_SEQS:=64}"
: "${BASE_GPU_UTIL:=0.7}"                           # KV-cache split between the two
: "${SMALL_GPU_UTIL:=0.1}"
export BASE_MODEL_NAME SMALL_MODEL_NAME BASE_MODEL_PORT SMALL_MODEL_PORT

# -------------------------------------------------------------------- steering
: "${STEER_VEC_PATH:=vector_500_500/DeepSeek-R1-Distill-Qwen-1.5B/layer_20_rank_60.pt}"
: "${STEER_SCALE:=1.0}"                             # alpha
: "${STATIC_STEER_LAYER:=20}"                       # layer the vLLM hook adds v at
: "${STEER_MATCH_TOKEN_IDS:=271}"                   # "\n\n" -> steer at step boundaries

# --------------------------------------------------------------------- routing
: "${WINDOW_SIZE:=10}"                              # W
: "${TT_RESHAPE:=16 16}"                            # d1 d2
: "${IMPL:=hf}"                                     # hidden-state backend: hf | vllm
: "${NUM_REPEATS:=1}"                               # runs per problem (pass@k -> k)
# Layer the hidden states are read from. The steering hook adds v at the *input*
# of STATIC_STEER_LAYER, so the steered state is the output of the layer below.
: "${HIDDEN_STATE_LAYER:=$((STATIC_STEER_LAYER - 1))}"

# ------------------------------------------------------------- sweep grids
# Overridable, e.g. RANK1_LIST="7" ENTROPY_LIST="1.1" bash scripts/sweep_grid_steer.sh
: "${RANK1_LIST:=6 7 8}"                            # T_r1
: "${RANK2_LIST:=60 70 80}"                      # T_r2
: "${ENTROPY_LIST:=0.9 1.1}"                        # T_e  (RankGuide + GlimpRouter)
: "${SPECREASON_LIST:=7 8 9}"                       # SpecReason score threshold
: "${EARLY_STOP:=1}"                                # 1 = enable collapse early-stop
: "${TASKS:=aime25}"                                # datasets to sweep
: "${TOKEN_BUDGET:=16384}"

# Which arms to run in a sweep script (1 = on, 0 = off)
: "${RUN_RANKGUIDE:=1}"
: "${RUN_GLIMPROUTER:=0}"
: "${RUN_SPECREASON:=0}"

# --------------------------------------------------------------------- timing
# Wall-clock of the shell driver, one CSV row per problem / config / sweep.
# (Per-step model latency lives inside the result pickles; see latency_breakdown.py.)
: "${TIMING_LOG:=${DATA_DIR}/rankguide/logs/timings.csv}"

# =============================================================== helper functions

# Seconds -> HH:MM:SS
fmt_hms() {
    printf '%02d:%02d:%02d' $(($1 / 3600)) $((($1 % 3600) / 60)) $(($1 % 60))
}

# Append one row to $TIMING_LOG.
#   log_timing <scope: problem|config|sweep> <label> <dataset> <seconds> [detail]
log_timing() {
    mkdir -p "$(dirname "$TIMING_LOG")"
    if [[ ! -s "$TIMING_LOG" ]]; then
        echo "timestamp,scope,label,dataset,seconds,hms,detail" >> "$TIMING_LOG"
    fi
    printf '%s,%s,%s,%s,%d,%s,%s\n' \
        "$(date +%Y-%m-%dT%H:%M:%S)" "$1" "$2" "$3" "$4" "$(fmt_hms "$4")" "${5:-}" \
        >> "$TIMING_LOG"
}

# Block until a vLLM server answers on $1.
wait_for_server() {
    local port=$1
    until curl -s "http://localhost:${port}/v1/models" > /dev/null; do
        echo "Waiting for server on port ${port} to start..."
        sleep 10
    done
    echo "Server on port ${port} is ready!"
}

# Launch the LRM, then the SRM. Pass "steer" as $1 to enable the static
# activation-steering hook on the SRM. Both servers are killed when the calling
# script exits unless KILL_SERVERS_ON_EXIT=0.
launch_servers() {
    local mode="${1:-nosteer}"

    python -c "import vllm; print('using vllm from', vllm.__file__)"

    vllm serve "$BASE_MODEL_NAME" --dtype auto -tp "$TP_SIZE" \
        --max_model_len "$MAX_MODEL_LEN" --gpu-memory-utilization "$BASE_GPU_UTIL" \
        --enable-prefix-caching --port "$BASE_MODEL_PORT" --max-num-seqs "$MAX_NUM_SEQS" &
    VLLM_BASE_PID=$!
    wait_for_server "$BASE_MODEL_PORT"

    if [[ "$mode" == "steer" ]]; then
        echo "SRM steering: ${STEER_VEC_PATH} @ layer ${STATIC_STEER_LAYER}, alpha=${STEER_SCALE}"
        STATIC_STEER_ENABLE=1 \
        STATIC_STEER_PATH="$STEER_VEC_PATH" \
        STATIC_STEER_SCALE="$STEER_SCALE" \
        STATIC_STEER_LAYER="$STATIC_STEER_LAYER" \
        STATIC_STEER_MATCH_TOKEN_IDS="$STEER_MATCH_TOKEN_IDS" \
        vllm serve "$SMALL_MODEL_NAME" --dtype auto -tp "$TP_SIZE" \
            --max_model_len "$MAX_MODEL_LEN" --gpu-memory-utilization "$SMALL_GPU_UTIL" \
            --enable-prefix-caching --port "$SMALL_MODEL_PORT" --max-num-seqs "$MAX_NUM_SEQS" &
    else
        vllm serve "$SMALL_MODEL_NAME" --dtype auto -tp "$TP_SIZE" \
            --max_model_len "$MAX_MODEL_LEN" --gpu-memory-utilization "$SMALL_GPU_UTIL" \
            --enable-prefix-caching --port "$SMALL_MODEL_PORT" --max-num-seqs "$MAX_NUM_SEQS" &
    fi
    VLLM_SMALL_PID=$!
    wait_for_server "$SMALL_MODEL_PORT"
    nvidia-smi

    # Also fires on Ctrl-C, so an aborted sweep does not leave two servers
    # holding the GPUs. Set KILL_SERVERS_ON_EXIT=0 to keep them up.
    if [[ "${KILL_SERVERS_ON_EXIT:-1}" == "1" ]]; then
        trap 'echo "Shutting down vLLM servers..."; kill '"$VLLM_BASE_PID $VLLM_SMALL_PID"' 2>/dev/null; exit' INT TERM
        trap 'kill '"$VLLM_BASE_PID $VLLM_SMALL_PID"' 2>/dev/null' EXIT
    fi
}

# Problem ids to sweep for a dataset.
problem_ids() {
    case "$1" in
        math)   seq 0 99  ;;
        aime)   seq 60 89 ;;
        aime25) seq 0 29  ;;
        gpqa)   seq 0 197 ;;
        lcbv5)  seq 0 166 ;;
        lcbv6)  seq 0 174 ;;
        *) echo "Unknown dataset: $1" >&2; return 1 ;;
    esac
}

# Name of the result directory for one config. The sweep scripts write it and
# the eval scripts read it, so keep this the single source of truth.
#   config_name <scheme> <T_r1> <T_r2> <T_e> <steering 0|1> <early_stop 0|1>
config_name() {
    local name="${1}_${2}_${3}_${WINDOW_SIZE}_${4}_${IMPL}"
    if [[ "$5" == "1" ]]; then name="${name}_steering_${STEER_SCALE}"; fi
    if [[ "$6" == "1" ]]; then name="${name}_early_stop"; fi
    printf '%s' "$name"
}

# Full results path for one config on one dataset.
#   result_dir <config_name> <dataset> <token_budget>
result_dir() {
    printf '%s' "${DATA_DIR}/rankguide/results/${1}/${2}_${3}/${BASE_MODEL_ABBRV}_${SMALL_MODEL_ABBRV}"
}