214 lines
7.5 KiB
Bash
Executable File
214 lines
7.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Launch a dynamo-sglang benchmark job on the GB200 cluster via srt-slurm.
|
|
#
|
|
# Required environment variables (set by the GitHub Actions workflow):
|
|
# FRAMEWORK - must be "dynamo-sglang"
|
|
# MODEL - HuggingFace model ID (used as fallback if no local path)
|
|
# MODEL_PREFIX - short prefix: "dsr1"
|
|
# PRECISION - "fp8" or "fp4"
|
|
# ISL - input sequence length (e.g. "1024")
|
|
# OSL - output sequence length (e.g. "1024")
|
|
# CONFIG_FILE - path relative to srt-slurm repo root (e.g. recipes/gb200-fp8/1k1k/low-latency.yaml)
|
|
# RESULT_FILENAME - prefix for output JSON filenames
|
|
# RUNNER_NAME - GitHub Actions runner name (used to tag the Slurm job)
|
|
# SQUASH_FILE - path to pre-imported sglang enroot squash file on Lustre
|
|
# NGINX_SQUASH_FILE - path to pre-imported nginx enroot squash file on Lustre
|
|
# SLURM_PARTITION - Slurm partition (default: batch)
|
|
# SLURM_ACCOUNT - Slurm account (default: sglang)
|
|
# SRT_SLURM_BRANCH - branch of srt-slurm repo to check out
|
|
# GITHUB_WORKSPACE - set automatically by GitHub Actions
|
|
|
|
set -euo pipefail
|
|
set -x
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Validate required vars
|
|
# ---------------------------------------------------------------------------
|
|
: "${FRAMEWORK:?}"
|
|
: "${MODEL_PREFIX:?}"
|
|
: "${PRECISION:?}"
|
|
: "${ISL:?}"
|
|
: "${OSL:?}"
|
|
: "${CONFIG_FILE:?}"
|
|
: "${RESULT_FILENAME:?}"
|
|
: "${RUNNER_NAME:?}"
|
|
: "${SQUASH_FILE:?}"
|
|
: "${NGINX_SQUASH_FILE:?}"
|
|
: "${GITHUB_WORKSPACE:?}"
|
|
|
|
SLURM_PARTITION="${SLURM_PARTITION:-batch}"
|
|
SLURM_ACCOUNT="${SLURM_ACCOUNT:-sglang}"
|
|
SRT_SLURM_BRANCH="${SRT_SLURM_BRANCH:-sglang-nightly-regression}"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Resolve local model paths on Lustre (avoids re-downloading on each run)
|
|
# ---------------------------------------------------------------------------
|
|
if [[ "$MODEL_PREFIX" == "dsr1" && "$PRECISION" == "fp8" ]]; then
|
|
MODEL_PATH="/mnt/lustre01/models/deepseek-r1-0528"
|
|
SRT_SLURM_MODEL_PREFIX="dsr1-fp8"
|
|
elif [[ "$MODEL_PREFIX" == "dsr1" && "$PRECISION" == "fp4" ]]; then
|
|
MODEL_PATH="/mnt/lustre01/models/deepseek-r1-0528-fp4-v2/"
|
|
SRT_SLURM_MODEL_PREFIX="dsr1-fp4"
|
|
else
|
|
MODEL_PATH="$MODEL"
|
|
SRT_SLURM_MODEL_PREFIX="$MODEL_PREFIX"
|
|
fi
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Set up per-runner Lustre workspace (cleaned before each run, accessible
|
|
# to both the runner and compute nodes)
|
|
# ---------------------------------------------------------------------------
|
|
LUSTRE_WORKSPACE="/mnt/lustre01/users-public/sglang-ci/workspace/${RUNNER_NAME}"
|
|
rm -rf "$LUSTRE_WORKSPACE"
|
|
mkdir -p "$LUSTRE_WORKSPACE"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Clone and set up srt-slurm
|
|
# ---------------------------------------------------------------------------
|
|
SRT_REPO_DIR="$LUSTRE_WORKSPACE/srt-slurm"
|
|
|
|
git clone https://github.com/NVIDIA/srt-slurm.git "$SRT_REPO_DIR"
|
|
cd "$SRT_REPO_DIR"
|
|
git checkout "$SRT_SLURM_BRANCH"
|
|
echo "--- srt-slurm last commit ---"
|
|
git log -1 --format="commit %H%nauthor %an%ndate %ad%nsubject %s" --date=iso
|
|
|
|
curl -LsSf https://astral.sh/uv/install.sh | sh
|
|
source "$HOME/.local/bin/env"
|
|
|
|
uv venv
|
|
source .venv/bin/activate
|
|
uv pip install -e .
|
|
|
|
if ! command -v srtctl &>/dev/null; then
|
|
echo "ERROR: srtctl installation failed"
|
|
exit 1
|
|
fi
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generate srtslurm.yaml
|
|
# ---------------------------------------------------------------------------
|
|
SRTCTL_ROOT="$SRT_REPO_DIR"
|
|
|
|
cat > srtslurm.yaml <<EOF
|
|
# SRT SLURM configuration for SGLang GB200 nightly CI
|
|
default_account: "${SLURM_ACCOUNT}"
|
|
default_partition: "${SLURM_PARTITION}"
|
|
default_time_limit: "6:00:00"
|
|
|
|
gpus_per_node: 4
|
|
network_interface: ""
|
|
|
|
srtctl_root: "${SRTCTL_ROOT}"
|
|
|
|
model_paths:
|
|
"${SRT_SLURM_MODEL_PREFIX}": "${MODEL_PATH}"
|
|
|
|
containers:
|
|
dynamo-sglang: ${SQUASH_FILE}
|
|
nginx: ${NGINX_SQUASH_FILE}
|
|
nginx-sqsh: ${NGINX_SQUASH_FILE}
|
|
EOF
|
|
|
|
echo "--- srtslurm.yaml ---"
|
|
cat srtslurm.yaml
|
|
|
|
make setup ARCH=aarch64
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Patch job name and submit via srtctl
|
|
# ---------------------------------------------------------------------------
|
|
sed -i "s/^name:.*/name: \"${RUNNER_NAME}\"/" "$CONFIG_FILE"
|
|
|
|
SRTCTL_OUTPUT=$(srtctl apply -f "$CONFIG_FILE" \
|
|
--tags "gb200,${MODEL_PREFIX},${PRECISION},${ISL}x${OSL},sglang-nightly-$(date +%Y%m%d)" \
|
|
--setup-script install-torchao.sh 2>&1)
|
|
echo "$SRTCTL_OUTPUT"
|
|
|
|
JOB_ID=$(echo "$SRTCTL_OUTPUT" | grep -oP '✅ Job \K[0-9]+' || echo "$SRTCTL_OUTPUT" | grep -oP 'Job \K[0-9]+' || true)
|
|
|
|
if [ -z "$JOB_ID" ]; then
|
|
echo "ERROR: Could not extract JOB_ID from srtctl output"
|
|
exit 1
|
|
fi
|
|
|
|
echo "Submitted Slurm job: $JOB_ID"
|
|
|
|
set +x
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Wait for job and stream logs
|
|
# ---------------------------------------------------------------------------
|
|
LOGS_DIR="outputs/$JOB_ID/logs"
|
|
LOG_FILE="$LOGS_DIR/sweep_${JOB_ID}.log"
|
|
|
|
mkdir -p "$LOGS_DIR"
|
|
|
|
while ! ls "$LOG_FILE" &>/dev/null; do
|
|
if ! squeue -j "$JOB_ID" --noheader 2>/dev/null | grep -q "$JOB_ID"; then
|
|
echo "ERROR: Job $JOB_ID failed before creating log file"
|
|
scontrol show job "$JOB_ID" || true
|
|
exit 1
|
|
fi
|
|
echo "Waiting for job $JOB_ID to start and $LOG_FILE to appear..."
|
|
sleep 5
|
|
done
|
|
|
|
(
|
|
while squeue -j "$JOB_ID" --noheader 2>/dev/null | grep -q "$JOB_ID"; do
|
|
sleep 10
|
|
done
|
|
) &
|
|
POLL_PID=$!
|
|
|
|
tail -F -s 2 -n+1 "$LOG_FILE" --pid=$POLL_PID 2>/dev/null
|
|
|
|
wait $POLL_PID
|
|
|
|
set -x
|
|
|
|
echo "Job $JOB_ID completed. Collecting results..."
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Collect results
|
|
# ---------------------------------------------------------------------------
|
|
if [ ! -d "$LOGS_DIR" ]; then
|
|
echo "WARNING: Logs directory not found at $LOGS_DIR"
|
|
exit 1
|
|
fi
|
|
|
|
cp -r "$LOGS_DIR" "$GITHUB_WORKSPACE/LOGS"
|
|
tar czf "$GITHUB_WORKSPACE/multinode_server_logs.tar.gz" -C "$LOGS_DIR" .
|
|
|
|
RESULT_SUBDIRS=$(find "$LOGS_DIR" -maxdepth 1 -type d -name "*isl*osl*" 2>/dev/null || true)
|
|
|
|
if [ -z "$RESULT_SUBDIRS" ]; then
|
|
echo "ERROR: No result subdirectories found in $LOGS_DIR — benchmark did not produce any output"
|
|
exit 1
|
|
else
|
|
RESULT_COUNT=0
|
|
for result_subdir in $RESULT_SUBDIRS; do
|
|
CONFIG_NAME=$(basename "$result_subdir")
|
|
RESULT_FILES=$(find "$result_subdir" -name "results_concurrency_*.json" 2>/dev/null || true)
|
|
for result_file in $RESULT_FILES; do
|
|
if [ -f "$result_file" ]; then
|
|
filename=$(basename "$result_file")
|
|
concurrency=$(echo "$filename" | sed -n 's/results_concurrency_\([0-9]*\)_gpus_.*/\1/p')
|
|
gpus=$(echo "$filename" | sed -n 's/results_concurrency_[0-9]*_gpus_\([0-9]*\)_ctx_.*/\1/p')
|
|
ctx=$(echo "$filename" | sed -n 's/.*_ctx_\([0-9]*\)_gen_.*/\1/p')
|
|
gen=$(echo "$filename" | sed -n 's/.*_gen_\([0-9]*\)\.json/\1/p')
|
|
DEST="$GITHUB_WORKSPACE/${RESULT_FILENAME}_${CONFIG_NAME}_conc${concurrency}_gpus_${gpus}_ctx_${ctx}_gen_${gen}.json"
|
|
cp "$result_file" "$DEST"
|
|
echo "Saved: $DEST"
|
|
RESULT_COUNT=$((RESULT_COUNT + 1))
|
|
fi
|
|
done
|
|
done
|
|
if [ "$RESULT_COUNT" -eq 0 ]; then
|
|
echo "ERROR: Result subdirectories found but no result JSON files produced — benchmark failed"
|
|
exit 1
|
|
fi
|
|
fi
|
|
|
|
echo "Done."
|