AMD PD/D PR ci (#17183)

Co-authored-by: YC Tseng <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>
Co-authored-by: bingxche <Bingxu.Chen@amd.com>
This commit is contained in:
Zhaoyi Li
2026-02-03 01:29:14 -06:00
committed by GitHub
parent 25508d11c0
commit 8e933e1914
5 changed files with 1071 additions and 0 deletions

View File

@@ -734,6 +734,118 @@ jobs:
run: |
bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 1 --timeout-per-file 3600
# =============================================== Disaggregation ====================================================
stage-b-test-large-8-gpu-35x-disaggregation-amd:
needs: [check-changes, stage-a-test-1-amd]
if: |
always() &&
(
(inputs.target_stage == 'stage-b-test-large-8-gpu-disaggregation-amd') ||
(
!inputs.target_stage &&
(!failure() && !cancelled()) &&
((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
)
)
strategy:
fail-fast: false
matrix:
runner: [linux-mi35x-gpu-8.fabric]
runs-on: ${{matrix.runner}}
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Ensure VRAM is clear
run: bash scripts/ensure_vram_clear.sh rocm
- name: Check Host RDMA Environment
id: rdma_detect
run: |
set +e
echo "=== Checking Host RDMA Environment ==="
echo ""
echo "=== 1. Ionic driver library check ==="
ls -l /usr/lib/x86_64-linux-gnu/libibverbs/libionic* 2>/dev/null || echo "libionic not found in standard path"
echo ""
echo "=== 2. Infiniband devices ==="
ls -la /dev/infiniband/ 2>/dev/null || echo "/dev/infiniband not found"
ls -la /sys/class/infiniband/ 2>/dev/null || echo "/sys/class/infiniband not found"
echo ""
echo "=== 3. ibv_devinfo ==="
which ibv_devinfo 2>/dev/null && ibv_devinfo 2>&1 || echo "ibv_devinfo not available"
echo ""
echo "=== 4. Kernel modules ==="
lsmod 2>/dev/null | grep -E "ib_|rdma|ionic" || echo "No RDMA kernel modules loaded"
echo ""
echo "=== 5. Detect RDMA Devices for test environment ==="
if [ -d "/sys/class/infiniband" ]; then
RDMA_DEVS=$(ls /sys/class/infiniband | paste -sd "," -)
echo "Detected RDMA Devices: $RDMA_DEVS"
echo "SGLANG_TEST_RDMA_DEVICE=$RDMA_DEVS" >> $GITHUB_ENV
else
echo "No RDMA devices found in /sys/class/infiniband"
echo "SGLANG_TEST_RDMA_DEVICE=" >> $GITHUB_ENV
fi
echo ""
echo "=== Host RDMA Check Complete ==="
- name: Start Special Container
run: bash scripts/ci/amd/amd_ci_start_container_disagg.sh
env:
GITHUB_WORKSPACE: ${{ github.workspace }}
- name: Install dependencies
run: bash scripts/ci/amd/amd_ci_install_dependency.sh
- name: Verify RDMA in Container
run: |
docker exec -u root ci_sglang bash -c '
echo "=== Container RDMA Verification ==="
echo "Device nodes:"
ls -la /dev/infiniband/
echo ""
echo "Provider libraries:"
ls /usr/lib/x86_64-linux-gnu/libibverbs/ | grep -E "ionic|mlx" || echo "No Ionic/Mellanox providers"
echo ""
echo "HCA devices:"
HCA_COUNT=$(ibv_devinfo -list 2>&1 | grep -oE "^[0-9]+ HCAs? found" | grep -oE "^[0-9]+" || echo "0")
ibv_devinfo -list
if [ "$HCA_COUNT" -gt 0 ]; then
echo ""
echo "=== SUCCESS: RDMA setup complete. Found $HCA_COUNT HCA(s) ==="
else
echo ""
echo "=== WARNING: No HCAs detected. RDMA tests may fail ==="
fi
'
- name: Run Aiter Op Test (RMSNorm)
timeout-minutes: 10
run: |
echo "Running pre-check: test_rmsnorm2d.py"
docker exec \
-e MAX_JOBS=192 \
ci_sglang \
python /sgl-workspace/aiter/op_tests/test_rmsnorm2d.py
- name: Run test_disaggregation
timeout-minutes: 60
run: |
bash scripts/ci/amd/amd_ci_exec.sh \
-e SGLANG_TEST_RDMA_DEVICE="${{ env.SGLANG_TEST_RDMA_DEVICE }}" \
-w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-8-gpu-35x-disaggregation-amd --timeout-per-file 1800
pr-test-amd-finish:
needs:
[
@@ -750,6 +862,7 @@ jobs:
stage-b-test-small-1-gpu-amd-mi35x,
stage-b-test-large-1-gpu-amd,
stage-b-test-large-2-gpu-amd,
stage-b-test-large-8-gpu-35x-disaggregation-amd,
stage-c-test-large-8-gpu-amd,
stage-c-test-large-8-gpu-amd-mi35x,
]

View File

@@ -0,0 +1,244 @@
#!/bin/bash
set -euo pipefail
# Get version from git tags
SGLANG_VERSION="v0.5.5" # Default version, will be overridden if git tags are found
# Fetch tags from origin to ensure we have the latest
if git fetch --tags origin; then
# Get the latest version tag sorted by version number (e.g., v0.5.7)
VERSION_FROM_TAG=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1)
if [ -n "$VERSION_FROM_TAG" ]; then
SGLANG_VERSION="$VERSION_FROM_TAG"
echo "Using SGLang version from git tags: $SGLANG_VERSION"
else
echo "Warning: No version tags found; using default $SGLANG_VERSION" >&2
fi
else
echo "Warning: Failed to fetch tags from origin; using default $SGLANG_VERSION" >&2
fi
# Default base tags (can be overridden by command line arguments)
ROCM_VERSION="rocm700"
DEFAULT_MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
DEFAULT_MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
# Parse command line arguments
MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
while [[ $# -gt 0 ]]; do
case $1 in
--mi30x-base-tag) MI30X_BASE_TAG="$2"; shift 2;;
--mi35x-base-tag) MI35X_BASE_TAG="$2"; shift 2;;
-h|--help)
echo "Usage: $0 [--mi30x-base-tag TAG] [--mi35x-base-tag TAG]"
exit 0
;;
*) echo "Unknown option $1"; exit 1;;
esac
done
# Detect GPU architecture from the Kubernetes runner hostname
HOSTNAME_VALUE=$(hostname)
GPU_ARCH="mi30x" # default
# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
GPU_ARCH="${BASH_REMATCH[1]}"
echo "Detected GPU architecture from hostname: ${GPU_ARCH}"
else
echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
fi
# Normalise / collapse architectures we dont yet build specifically for
case "${GPU_ARCH}" in
mi35x)
echo "Runner uses ${GPU_ARCH}; will fetch mi35x image."
;;
mi30x|mi300|mi325)
echo "Runner uses ${GPU_ARCH}; will fetch mi30x image."
GPU_ARCH="mi30x"
;;
*)
echo "Runner architecture '${GPU_ARCH}' unrecognised; defaulting to mi30x image." >&2
GPU_ARCH="mi30x"
;;
esac
# Set up DEVICE_FLAG based on Kubernetes pod info
if [[ -f /etc/podinfo/gha-render-devices ]]; then
DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
else
DEVICE_FLAG="--device /dev/dri"
fi
# Find the latest image
find_latest_image() {
local gpu_arch=$1
local base_tag days_back image_tag
case "${gpu_arch}" in
mi30x) base_tag="${MI30X_BASE_TAG}" ;;
mi35x) base_tag="${MI35X_BASE_TAG}" ;;
*) echo "Error: unsupported GPU architecture '${gpu_arch}'" >&2; return 1 ;;
esac
# First, check local cache
for days_back in {0..6}; do
image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
local local_image="rocm/sgl-dev:${image_tag}"
image_id=$(docker images -q "${local_image}")
if [[ -n "$image_id" ]]; then
echo "Found cached image locally: ${local_image}" >&2
echo "${local_image}"
return 0
fi
done
# If not found locally, fall back to pulling from public registry
for days_back in {0..6}; do
image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
echo "Checking for image: rocm/sgl-dev:${image_tag}" >&2
if docker manifest inspect "rocm/sgl-dev:${image_tag}" >/dev/null 2>&1; then
echo "Found available image: rocm/sgl-dev:${image_tag}" >&2
echo "rocm/sgl-dev:${image_tag}"
return 0
fi
done
# If still not found, try finding any image matching ROCm+arch from remote registry
echo "Exact version not found. Searching remote registry for any ${ROCM_VERSION}-${gpu_arch} image…" >&2
for days_back in {0..6}; do
local target_date=$(date -d "${days_back} days ago" +%Y%m%d)
local remote_tags=$(curl -s "https://registry.hub.docker.com/v2/repositories/rocm/sgl-dev/tags?page_size=100&name=${ROCM_VERSION}-${gpu_arch}-${target_date}" 2>/dev/null | grep -o '"name":"[^"]*"' | cut -d'"' -f4 | head -n 1)
if [[ -n "$remote_tags" ]]; then
echo "Found available image: rocm/sgl-dev:${remote_tags}" >&2
echo "rocm/sgl-dev:${remote_tags}"
return 0
fi
done
echo "No recent images found. Searching any cached local images matching ROCm+arch…" >&2
local any_local
any_local=$(docker images --format '{{.Repository}}:{{.Tag}}' --filter "reference=rocm/sgl-dev:*${ROCM_VERSION}*${gpu_arch}*" | sort -r | head -n 1)
if [[ -n "$any_local" ]]; then
echo "Using cached fallback image: ${any_local}" >&2
echo "${any_local}"
return 0
fi
echo "Error: no ${gpu_arch} image found in the last 7 days for base ${base_tag}" >&2
echo "Using hard-coded fallback…" >&2
if [[ "${gpu_arch}" == "mi35x" ]]; then
echo "rocm/sgl-dev:v0.5.5-rocm700-mi35x-20251110"
else
echo "rocm/sgl-dev:v0.5.5-rocm700-mi30x-20251110"
fi
}
# Pull and run the latest image
IMAGE=$(find_latest_image "${GPU_ARCH}")
echo "Pulling Docker image: ${IMAGE}"
docker pull "${IMAGE}"
CACHE_HOST=/home/runner/sgl-data
if [[ -d "$CACHE_HOST" ]]; then
CACHE_VOLUME="-v $CACHE_HOST:/sgl-data"
else
CACHE_VOLUME=""
fi
# Detect libionic library for RDMA support
LIBIONIC_MOUNT=""
IONIC_SYMLINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
if [[ -L "$IONIC_SYMLINK" ]]; then
LIBIONIC_LIB=$(readlink -f "$IONIC_SYMLINK" 2>/dev/null)
if [[ -f "$LIBIONIC_LIB" ]]; then
echo "Found libionic library: $LIBIONIC_LIB (resolved from symlink)"
LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
else
echo "Warning: libionic symlink exists but target does not: $LIBIONIC_LIB"
fi
else
# Fallback: try to find directly
LIBIONIC_FOUND=$(find /usr/lib/x86_64-linux-gnu -maxdepth 1 -name "libionic.so.*" 2>/dev/null | head -1)
if [[ -n "$LIBIONIC_FOUND" ]]; then
LIBIONIC_LIB=$(readlink -f "$LIBIONIC_FOUND" 2>/dev/null)
if [[ -f "$LIBIONIC_LIB" ]]; then
echo "Found libionic library: $LIBIONIC_LIB"
LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
else
echo "Warning: libionic found but cannot resolve real path: $LIBIONIC_FOUND"
fi
else
echo "Warning: libionic library not found on host, RDMA may not work"
fi
fi
MOUNT_ARGS=""
add_mount_if_exists() {
local name=$1
local search_pattern=$2
local path=$(find /lib/x86_64-linux-gnu /usr/lib/x86_64-linux-gnu /lib64 /usr/lib64 -name "$search_pattern" -print -quit 2>/dev/null)
if [ -n "$path" ]; then
echo "Found $name at: $path"
MOUNT_ARGS="$MOUNT_ARGS -v $path:$path:ro"
else
echo "WARNING: Could not find $name on host! (Pattern: $search_pattern)"
fi
}
IONIC_LINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
if [ -L "$IONIC_LINK" ]; then
IONIC_REAL=$(readlink -f "$IONIC_LINK")
if [ -f "$IONIC_REAL" ]; then
echo "Ionic Driver: $IONIC_REAL"
MOUNT_ARGS="$MOUNT_ARGS -v $IONIC_REAL:$IONIC_REAL:ro"
fi
fi
add_mount_if_exists "libnl-3" "libnl-3.so*"
add_mount_if_exists "libmnl" "libmnl.so*"
echo "Mount args: $MOUNT_ARGS"
echo "Launching container: ci_sglang"
docker run -dt --user root \
--device=/dev/kfd \
--device=/dev/dri \
${DEVICE_FLAG} \
-v "${GITHUB_WORKSPACE:-$PWD}:/sglang-checkout" \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
-v /sys/class/infiniband_verbs:/sys/class/infiniband_verbs:ro \
-v /sys/class/net:/sys/class/net:ro \
-v /etc/libibverbs.d:/etc/libibverbs.d:ro \
-v /usr/lib/x86_64-linux-gnu/libibverbs:/usr/lib/x86_64-linux-gnu/libibverbs:ro \
$MOUNT_ARGS \
$CACHE_VOLUME \
--privileged \
--network=host \
--ipc=host \
--ulimit memlock=-1 \
--cap-add=IPC_LOCK \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--group-add video \
--group-add rdma \
--shm-size 32g \
-e HF_TOKEN="${HF_TOKEN:-}" \
-e HF_HOME=/sgl-data/hf-cache \
-e HF_HUB_ETAG_TIMEOUT=300 \
-e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
-e MIOPEN_CUSTOM_CACHE_DIR=/sgl-data/miopen-cache \
-w /sglang-checkout \
--name ci_sglang \
"${IMAGE}"

View File

@@ -0,0 +1,422 @@
import json
import os
import unittest
from types import SimpleNamespace
import openai
import requests
from transformers import AutoTokenizer
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
register_amd_ci(est_time=600, suite="stage-b-test-large-8-gpu-35x-disaggregation-amd")
class TestDisaggregationAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Configure ROCm RDMA environment
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--attention-backend",
"aiter",
"--log-level",
"debug",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.8",
"--log-level",
"debug",
]
decode_args += cls.transfer_backend + cls.rdma_devices
print("Debug")
print(decode_args)
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.70)
def test_logprob(self):
prompt = "The capital of france is "
response = requests.post(
self.lb_url + "/generate",
json={
"text": prompt,
"sampling_params": {"temperature": 0},
"return_logprob": True,
"return_input_logprob": True,
"logprob_start_len": 0,
},
)
j = response.json()
completion_tokens = j["meta_info"]["completion_tokens"]
input_logprobs = j["meta_info"]["input_token_logprobs"]
output_logprobs = j["meta_info"]["output_token_logprobs"]
assert (
len(output_logprobs) == completion_tokens
), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}"
assert (
len(input_logprobs) > 0
), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
def test_structured_output(self):
json_schema = json.dumps(
{
"type": "object",
"properties": {
"name": {"type": "string", "pattern": "^[\\w]+$"},
"population": {"type": "integer"},
},
"required": ["name", "population"],
}
)
# JSON
response = requests.post(
f"{self.lb_url}/generate",
json={
"text": "Here is the information of the capital of France in the JSON format.\n",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 64,
"json_schema": json_schema,
},
},
)
output = response.json()["text"]
# ensure the output is a valid JSON
json.loads(output)
def test_first_token_finish(self):
client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
tokenizer = AutoTokenizer.from_pretrained(self.model)
eos_token = tokenizer.eos_token_id
prompt = "The best programming language for AI is"
# First token EOS
res = client.completions.create(
model="dummy", prompt=prompt, logit_bias={eos_token: 42}
).model_dump()
print(f"{res=}")
assert res["usage"]["completion_tokens"] == 1, (
"Expected completion_tokens to be 1 when first token is EOS, "
f"but got {res['usage']['completion_tokens']}"
)
# First token EOS with ignore_eos
res = client.completions.create(
model="dummy",
prompt=prompt,
logit_bias={eos_token: 42},
extra_body={"ignore_eos": True},
).model_dump()
print(f"{res=}")
assert res["usage"]["completion_tokens"] > 1, (
"Expected completion_tokens to be greater than 1 when ignore_eos is True, "
f"but got {res['usage']['completion_tokens']}"
)
# First token with specified stop token
stop_token_id = tokenizer.encode(" hello", add_special_tokens=False)[0]
res = client.completions.create(
model="dummy",
prompt=prompt,
logit_bias={stop_token_id: 42},
stop=[" hello"],
).model_dump()
print(f"{res=}")
assert res["usage"]["completion_tokens"] == 1, (
"Expected completion_tokens to be 1 when first token is stop token, "
f"but got {res['usage']['completion_tokens']}"
)
# register_amd_ci(est_time=300, suite="stage-b-test-large-2-gpu-amd")
class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Configure ROCm RDMA environment
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
# set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def tearDownClass(cls):
os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB")
super().tearDownClass()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--attention-backend",
"aiter",
"--log-level",
"debug",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.8",
"--log-level",
"debug",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
# Expect lots of failure but the server cannot crash
try:
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
except Exception as e:
print(f"Test encountered expected errors: {e}")
# Check if servers are still healthy
try:
response = requests.get(self.prefill_url + "/health_generate")
assert response.status_code == 200
response = requests.get(self.decode_url + "/health_generate")
assert response.status_code == 200
except Exception as health_check_error:
# If health check fails, re-raise the original exception
raise e from health_check_error
# register_amd_ci(est_time=300, suite="stage-b-test-large-2-gpu-amd")
class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# Configure ROCm RDMA environment
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
os.environ["SGLANG_TEST_RETRACT"] = "true"
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def tearDownClass(cls):
os.environ.pop("SGLANG_TEST_RETRACT")
super().tearDownClass()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--attention-backend",
"aiter",
"--log-level",
"debug",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.8",
"--log-level",
"debug",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"Evaluation metrics: {metrics}")
self.assertGreater(metrics["accuracy"], 0.70)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,291 @@
import os
import time
import unittest
from types import SimpleNamespace
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
try_cached_model,
)
register_amd_ci(est_time=600, suite="stage-b-test-large-8-gpu-35x-disaggregation-amd")
class TestDisaggregationPrefillPPAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# set up ROCm env
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
"--attention-backend",
"aiter",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--base-gpu-id",
"4",
"--attention-backend",
"aiter",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.70)
# Wait a little bit so that the memory check happens.
time.sleep(5)
# register_amd_ci(est_time=200, suite="stage-c-test-large-8-gpu-amd")
class TestDisaggregationPrefillPPDynamicChunkAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# set up ROCm env
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
"--enable-dynamic-chunking",
"--attention-backend",
"aiter",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--base-gpu-id",
"4",
"--attention-backend",
"aiter",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.70)
# Wait a little bit so that the memory check happens.
time.sleep(5)
# register_amd_ci(est_time=200, suite="stage-c-test-large-8-gpu-amd")
class TestDisaggregationDecodePPAccuracy(PDDisaggregationServerBase):
@classmethod
def setUpClass(cls):
super().setUpClass()
# set up ROCm env
os.environ["SGLANG_USE_AITER"] = "1"
rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE")
if rdma_env:
cls.rdma_devices = ["--disaggregation-ib-device", rdma_env]
print(f"Found RDMA devices in env: {rdma_env}")
else:
print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
cls.rdma_devices = []
cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
# Non blocking start servers
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"2",
"--pp-size",
"2",
"--disable-overlap-schedule",
"--attention-backend",
"aiter",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"2",
"--pp-size",
"2",
"--base-gpu-id",
"4",
"--attention-backend",
"aiter",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host=f"http://{self.base_host}",
port=int(self.lb_port),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.70)
# Wait a little bit so that the memory check happens.
time.sleep(5)
if __name__ == "__main__":
unittest.main()

View File

@@ -22,6 +22,7 @@ PER_COMMIT_SUITES = {
"stage-a-test-1-amd",
"stage-b-test-small-1-gpu-amd",
"stage-b-test-small-1-gpu-amd-mi35x",
"stage-b-test-large-8-gpu-35x-disaggregation-amd",
"stage-b-test-large-1-gpu-amd",
"stage-b-test-large-2-gpu-amd",
"stage-c-test-large-8-gpu-amd-mi35x",