diff --git a/11_start_server_122b.sh b/11_start_server_122b.sh index 67b3b58..da71741 100755 --- a/11_start_server_122b.sh +++ b/11_start_server_122b.sh @@ -17,9 +17,22 @@ # GPU_MEM_UTIL - GPU memory utilization fraction (default: 0.92) # API_KEY - API key for authentication (default: none) # TENSOR_PARALLEL - Number of GPUs for tensor parallelism (default: 4) -# MAX_NUM_SEQS - Max concurrent sequences (default: 4). Raise this +# MAX_NUM_SEQS - Max concurrent sequences (default: 16). Raise this # when more users need it AND the KV cache fits; # lower it if 128k context runs out of memory. +# MAX_BATCHED_TOK - Token budget per scheduler step (default: 8192). +# With chunked prefill, long 128k prompts are split +# into chunks of this size so they don't block other +# users' decode steps. Key for concurrent throughput. +# KV_CACHE_DTYPE - KV cache quantization (default: fp8). Halves KV +# memory vs bf16 -> ~2x more concurrent 128k slots. +# Set to "auto" to disable. +# +# NOTE: For best concurrent throughput on these 4x L40S (no NVLink), +# prefer the TWO-REPLICA setup (TP=2 x2 + load balancer): +# bash 14_start_replicas_122b.sh +# TP=4 pays a per-layer all-reduce tax over PCIe; two TP=2 +# replicas avoid half of it and double the request slots. # ------------------------------------------------------------------ set -euo pipefail @@ -29,10 +42,12 @@ SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif" MODEL_DIR="${MODEL_DIR:-$HOME/models/Qwen3.5-122B-A10B-FP8}" PORT="${PORT:-7080}" MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}" -GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.92}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}" API_KEY="${API_KEY:-}" TENSOR_PARALLEL="${TENSOR_PARALLEL:-4}" -MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-16}" +MAX_BATCHED_TOK="${MAX_BATCHED_TOK:-8192}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" if [ ! -f "$SIF_FILE" ]; then echo "ERROR: Container image not found at ${SIF_FILE}" @@ -61,7 +76,10 @@ VLLM_ARGS=( --reasoning-parser qwen3 --served-model-name "qwen3.5-122b-a10b-fp8" --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_BATCHED_TOK" + --kv-cache-dtype "$KV_CACHE_DTYPE" --enable-prefix-caching + --enable-chunked-prefill ) if [ -n "$API_KEY" ]; then @@ -78,6 +96,8 @@ echo " Context len: ${MAX_MODEL_LEN}" echo " GPU util: ${GPU_MEM_UTIL}" echo " TP size: ${TENSOR_PARALLEL}" echo " Max seqs: ${MAX_NUM_SEQS}" +echo " Batched tok: ${MAX_BATCHED_TOK} (chunked prefill)" +echo " KV dtype: ${KV_CACHE_DTYPE}" echo " dtype: auto (FP8)" echo " API key: ${API_KEY:-}" echo "==============================================" diff --git a/14_start_replica_122b.sh b/14_start_replica_122b.sh new file mode 100755 index 0000000..beefb1f --- /dev/null +++ b/14_start_replica_122b.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +# ------------------------------------------------------------------ +# 14_start_replica_122b.sh +# Launches ONE vLLM replica of Qwen3.5-122B-A10B-FP8 with TP=2, +# pinned to a specific pair of GPUs, on its own internal port. +# +# Two of these (GPUs 0,1 and GPUs 2,3) behind the load balancer +# (16_start_loadbalancer.sh) beat a single TP=4 server on these +# L40S cards: TP=4 does a per-layer all-reduce over PCIe (no NVLink), +# which is pure overhead. Two TP=2 replicas halve that comm cost per +# token AND double the number of concurrent request slots. +# +# This is normally invoked by 15_start_replicas_122b.sh, not directly. +# +# Required environment variables: +# GPUS - GPU pair, e.g. "0,1" or "2,3" +# PORT - Internal port for this replica, e.g. 7081 / 7082 +# +# Optional (sane defaults below): +# MODEL_DIR, MAX_MODEL_LEN, GPU_MEM_UTIL, API_KEY, +# MAX_NUM_SEQS, MAX_BATCHED_TOK, KV_CACHE_DTYPE +# ------------------------------------------------------------------ +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif" + +: "${GPUS:?Set GPUS, e.g. GPUS=0,1}" +: "${PORT:?Set PORT, e.g. PORT=7081}" + +MODEL_DIR="${MODEL_DIR:-$HOME/models/Qwen3.5-122B-A10B-FP8}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.95}" +API_KEY="${API_KEY:-}" +# Per-replica concurrency. Two replicas => 2x this many total slots. +MAX_NUM_SEQS="${MAX_NUM_SEQS:-16}" +MAX_BATCHED_TOK="${MAX_BATCHED_TOK:-8192}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" + +if [ ! -f "$SIF_FILE" ]; then + echo "ERROR: Container image not found at ${SIF_FILE}" + echo " Run 01_build_container.sh first." + exit 1 +fi + +if [ ! -d "$MODEL_DIR" ]; then + echo "ERROR: Model directory not found at ${MODEL_DIR}" + echo " Run 10_download_model_122b.sh first." + exit 1 +fi + +MODEL_PARENT="$(dirname "$MODEL_DIR")" +MODEL_NAME="$(basename "$MODEL_DIR")" + +VLLM_ARGS=( + --model "/models/${MODEL_NAME}" + --port "$PORT" + --host 0.0.0.0 + --tensor-parallel-size 2 + --max-model-len "$MAX_MODEL_LEN" + --gpu-memory-utilization "$GPU_MEM_UTIL" + --dtype auto + --trust-remote-code + --reasoning-parser qwen3 + --served-model-name "qwen3.5-122b-a10b-fp8" + --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_BATCHED_TOK" + --kv-cache-dtype "$KV_CACHE_DTYPE" + --enable-prefix-caching + --enable-chunked-prefill +) + +if [ -n "$API_KEY" ]; then + VLLM_ARGS+=(--api-key "$API_KEY") +fi + +echo "==============================================" +echo " vLLM REPLICA — Qwen3.5-122B-A10B-FP8 (TP=2)" +echo "==============================================" +echo " GPUs: ${GPUS}" +echo " Port: ${PORT}" +echo " Context len: ${MAX_MODEL_LEN}" +echo " GPU util: ${GPU_MEM_UTIL}" +echo " Max seqs: ${MAX_NUM_SEQS}" +echo " Batched tok: ${MAX_BATCHED_TOK} (chunked prefill)" +echo " KV dtype: ${KV_CACHE_DTYPE}" +echo "==============================================" + +# CUDA_VISIBLE_DEVICES pins this replica to its GPU pair. vLLM then +# sees them as logical devices 0 and 1 for its TP=2. +apptainer exec --nv \ + --writable-tmpfs \ + --env CUDA_VISIBLE_DEVICES="${GPUS}" \ + --bind "${MODEL_PARENT}:/models" \ + "$SIF_FILE" \ + python3 -m vllm.entrypoints.openai.api_server "${VLLM_ARGS[@]}" diff --git a/15_start_replicas_122b.sh b/15_start_replicas_122b.sh new file mode 100755 index 0000000..7df6d99 --- /dev/null +++ b/15_start_replicas_122b.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# ------------------------------------------------------------------ +# 15_start_replicas_122b.sh +# Starts TWO vLLM replicas of Qwen3.5-122B-A10B-FP8 in the background: +# replica A: GPUs 0,1 -> internal port 7091 +# replica B: GPUs 2,3 -> internal port 7092 +# Then starts the load balancer on the public port 7080. +# +# Ports 7091/7092 are used (not 7081/7082) because Open WebUI runs on +# 7081 — the replicas must not collide with it. +# +# This replaces the single TP=4 server (11/12_*.sh) for better +# concurrent throughput on the 4x L40S box. Clients keep using +# http://:7080/v1 exactly as before. +# +# Usage: +# bash 15_start_replicas_122b.sh +# +# Stop everything: bash 17_stop_replicas_122b.sh +# ------------------------------------------------------------------ +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +LOG_DIR="${SCRIPT_DIR}/logs" +mkdir -p "$LOG_DIR" + +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" + +REPLICA_A_PORT="${REPLICA_A_PORT:-7091}" +REPLICA_B_PORT="${REPLICA_B_PORT:-7092}" +PUBLIC_PORT="${PUBLIC_PORT:-7080}" +export API_KEY="${API_KEY:-}" + +# Refuse to start if anything is already bound to the public port. +if [ -f "${LOG_DIR}/vllm_server.pid" ]; then + OLD_PID=$(cat "${LOG_DIR}/vllm_server.pid") + if kill -0 "$OLD_PID" 2>/dev/null; then + echo "A single-server vLLM (PID ${OLD_PID}) is still running." + echo "Stop it first: bash 05_stop_server.sh" + exit 1 + fi +fi + +start_replica () { + local name="$1" gpus="$2" port="$3" + local log="${LOG_DIR}/vllm_replica_${name}_${TIMESTAMP}.log" + local pidfile="${LOG_DIR}/vllm_replica_${name}.pid" + + if [ -f "$pidfile" ] && kill -0 "$(cat "$pidfile")" 2>/dev/null; then + echo "Replica ${name} already running (PID $(cat "$pidfile"))." + return + fi + + echo "Starting replica ${name}: GPUs ${gpus} -> port ${port}" + echo " log: ${log}" + GPUS="$gpus" PORT="$port" \ + nohup bash "${SCRIPT_DIR}/14_start_replica_122b.sh" > "$log" 2>&1 & + echo "$!" > "$pidfile" + echo " PID: $(cat "$pidfile")" +} + +start_replica "a" "0,1" "$REPLICA_A_PORT" +start_replica "b" "2,3" "$REPLICA_B_PORT" + +echo "" +echo "Both replicas launching (model load takes 5-10 min each, in parallel)." +echo "Watch: tail -f ${LOG_DIR}/vllm_replica_*_${TIMESTAMP}.log" +echo "" + +# Start the load balancer on the public port. +REPLICA_A_PORT="$REPLICA_A_PORT" REPLICA_B_PORT="$REPLICA_B_PORT" \ +PUBLIC_PORT="$PUBLIC_PORT" \ + bash "${SCRIPT_DIR}/16_start_loadbalancer.sh" + +echo "" +echo "Public API: http://$(hostname):${PUBLIC_PORT}/v1" +echo "Stop all: bash 17_stop_replicas_122b.sh" diff --git a/16_start_loadbalancer.sh b/16_start_loadbalancer.sh new file mode 100755 index 0000000..63b102d --- /dev/null +++ b/16_start_loadbalancer.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# ------------------------------------------------------------------ +# 16_start_loadbalancer.sh +# Starts the streaming load balancer (lb_proxy.py) on the public port, +# fanning requests out to the two 122B replicas by least-in-flight. +# +# Runs inside the vLLM container (Python + httpx/uvicorn already there). +# Normally invoked by 15_start_replicas_122b.sh, but can run alone. +# +# Env vars: +# PUBLIC_PORT public port (default 7080) +# REPLICA_A_PORT replica A port (default 7091) +# REPLICA_B_PORT replica B port (default 7092) +# ------------------------------------------------------------------ +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +SIF_FILE="${SCRIPT_DIR}/vllm_qwen.sif" +LOG_DIR="${SCRIPT_DIR}/logs" +mkdir -p "$LOG_DIR" + +PUBLIC_PORT="${PUBLIC_PORT:-7080}" +REPLICA_A_PORT="${REPLICA_A_PORT:-7091}" +REPLICA_B_PORT="${REPLICA_B_PORT:-7092}" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" +LOG_FILE="${LOG_DIR}/vllm_lb_${TIMESTAMP}.log" +PID_FILE="${LOG_DIR}/vllm_lb.pid" + +if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then + echo "Load balancer already running (PID $(cat "$PID_FILE"))." + exit 0 +fi + +BACKENDS="http://127.0.0.1:${REPLICA_A_PORT},http://127.0.0.1:${REPLICA_B_PORT}" + +echo "Starting load balancer on port ${PUBLIC_PORT} -> ${BACKENDS}" +echo " log: ${LOG_FILE}" + +# --net not needed: the proxy and replicas share the host network. +# Install starlette/uvicorn/httpx into a tmpfs if the image lacks them. +nohup apptainer exec --writable-tmpfs \ + --env PUBLIC_PORT="$PUBLIC_PORT" \ + --env BACKENDS="$BACKENDS" \ + --bind "${SCRIPT_DIR}:${SCRIPT_DIR}" \ + "$SIF_FILE" \ + bash -c "python3 -c 'import starlette,uvicorn,httpx' 2>/dev/null || pip install --quiet starlette uvicorn httpx; exec python3 '${SCRIPT_DIR}/lb_proxy.py'" \ + > "$LOG_FILE" 2>&1 & + +echo "$!" > "$PID_FILE" +echo " PID: $(cat "$PID_FILE")" diff --git a/17_stop_replicas_122b.sh b/17_stop_replicas_122b.sh new file mode 100755 index 0000000..45b53f0 --- /dev/null +++ b/17_stop_replicas_122b.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# ------------------------------------------------------------------ +# 17_stop_replicas_122b.sh +# Stops the load balancer and both 122B replicas started by +# 15_start_replicas_122b.sh. +# ------------------------------------------------------------------ +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +LOG_DIR="${SCRIPT_DIR}/logs" + +stop_one () { + local name="$1" pidfile="$2" + if [ ! -f "$pidfile" ]; then + echo "${name}: no PID file, skipping." + return + fi + local pid + pid=$(cat "$pidfile") + if kill -0 "$pid" 2>/dev/null; then + echo "Stopping ${name} (PID ${pid})..." + kill "$pid" 2>/dev/null || true + for _ in 1 2 3 4 5; do + kill -0 "$pid" 2>/dev/null || break + sleep 1 + done + if kill -0 "$pid" 2>/dev/null; then + echo " still alive, SIGKILL..." + kill -9 "$pid" 2>/dev/null || true + fi + echo " ${name} stopped." + else + echo "${name} (PID ${pid}) not running." + fi + rm -f "$pidfile" +} + +# Stop LB first so no new requests get routed to dying replicas. +stop_one "load balancer" "${LOG_DIR}/vllm_lb.pid" +stop_one "replica a" "${LOG_DIR}/vllm_replica_a.pid" +stop_one "replica b" "${LOG_DIR}/vllm_replica_b.pid" + +echo "" +echo "All replicas and load balancer stopped." +echo "Note: vLLM worker subprocesses may take a few seconds to release GPU memory." +echo "Verify with: nvidia-smi" diff --git a/lb_proxy.py b/lb_proxy.py new file mode 100755 index 0000000..ee52b45 --- /dev/null +++ b/lb_proxy.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +""" +lb_proxy.py — minimal streaming reverse proxy / load balancer for the +two vLLM 122B replicas. + +Why not nginx: this needs to stream Server-Sent Events (token-by-token +chat completions) without buffering, route only to replicas that are +actually up, and have zero install footprint. A ~120-line asyncio proxy +does exactly that and runs inside the existing vLLM container (which has +Python + the stdlib; we also use httpx, which ships in the vLLM image). + +Routing: least-busy of the healthy replicas. We track in-flight requests +per backend and send each new request to the one with the fewest. This +beats blind round-robin when one user fires a huge 128k prompt — the +other replica keeps serving short requests. + +Health: a backend that fails to connect is marked down and skipped; +we re-probe it on the next request after a short cooldown. + +Env vars: + PUBLIC_PORT public port to listen on (default 7080) + BACKENDS comma-separated backend base URLs + (default http://127.0.0.1:7081,http://127.0.0.1:7082) +""" +import asyncio +import os +import time + +import httpx +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import Response, StreamingResponse, JSONResponse +from starlette.routing import Route +import uvicorn + +BACKENDS = [ + b.strip() + for b in os.environ.get( + "BACKENDS", "http://127.0.0.1:7091,http://127.0.0.1:7092" + ).split(",") + if b.strip() +] +PUBLIC_PORT = int(os.environ.get("PUBLIC_PORT", "7080")) +COOLDOWN_S = 10.0 # how long to keep a backend marked down before retry + +# Per-backend state: in-flight count + "down until" timestamp. +inflight = {b: 0 for b in BACKENDS} +down_until = {b: 0.0 for b in BACKENDS} + +# Long timeout: 128k-context generations can run for minutes. +client = httpx.AsyncClient(timeout=httpx.Timeout(None, connect=10.0)) + + +def pick_backend() -> str | None: + now = time.monotonic() + healthy = [b for b in BACKENDS if down_until[b] <= now] + if not healthy: + # All cooling down — try the soonest-available one anyway. + healthy = BACKENDS + return min(healthy, key=lambda b: inflight[b]) + + +async def proxy(request: Request) -> Response: + backend = pick_backend() + if backend is None: + return JSONResponse({"error": "no backend available"}, status_code=503) + + url = backend + request.url.path + if request.url.query: + url += "?" + request.url.query + + # Drop hop-by-hop / host headers; pass the rest (incl. Authorization). + headers = { + k: v + for k, v in request.headers.items() + if k.lower() not in ("host", "content-length", "connection") + } + body = await request.body() + + inflight[backend] += 1 + try: + req = client.build_request( + request.method, url, headers=headers, content=body + ) + upstream = await client.send(req, stream=True) + except (httpx.ConnectError, httpx.ConnectTimeout): + down_until[backend] = time.monotonic() + COOLDOWN_S + inflight[backend] -= 1 + # Retry once on the other backend. + other = pick_backend() + if other and other != backend: + return await proxy(request) + return JSONResponse({"error": "backend unreachable"}, status_code=502) + + resp_headers = { + k: v + for k, v in upstream.headers.items() + if k.lower() not in ("content-length", "transfer-encoding", "connection") + } + + async def stream(): + try: + async for chunk in upstream.aiter_raw(): + yield chunk + finally: + await upstream.aclose() + inflight[backend] -= 1 + + return StreamingResponse( + stream(), + status_code=upstream.status_code, + headers=resp_headers, + media_type=upstream.headers.get("content-type"), + ) + + +async def health(request: Request) -> Response: + now = time.monotonic() + return JSONResponse( + { + "backends": [ + { + "url": b, + "inflight": inflight[b], + "up": down_until[b] <= now, + } + for b in BACKENDS + ] + } + ) + + +app = Starlette( + routes=[ + Route("/_lb_health", health, methods=["GET"]), + Route("/{path:path}", proxy, methods=["GET", "POST", "PUT", "DELETE"]), + ] +) + +if __name__ == "__main__": + print(f"Load balancer on :{PUBLIC_PORT} -> {BACKENDS}", flush=True) + uvicorn.run(app, host="0.0.0.0", port=PUBLIC_PORT, log_level="warning")