From 8e933e1914b6b335e58bb1442b3931e2303192bb Mon Sep 17 00:00:00 2001 From: Zhaoyi Li <36555117+Lzy17@users.noreply.github.com> Date: Tue, 3 Feb 2026 01:29:14 -0600 Subject: [PATCH] AMD PD/D PR ci (#17183) Co-authored-by: YC Tseng Co-authored-by: Bingxu Chen Co-authored-by: bingxche --- .github/workflows/pr-test-amd.yml | 113 +++++ .../ci/amd/amd_ci_start_container_disagg.sh | 244 ++++++++++ .../test_disaggregation_basic.py | 422 ++++++++++++++++++ .../disaggregation/test_disaggregation_pp.py | 291 ++++++++++++ test/run_suite.py | 1 + 5 files changed, 1071 insertions(+) create mode 100755 scripts/ci/amd/amd_ci_start_container_disagg.sh create mode 100644 test/registered/amd/disaggregation/test_disaggregation_basic.py create mode 100644 test/registered/amd/disaggregation/test_disaggregation_pp.py diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 769775cee..c0a6cd56b 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -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, ] diff --git a/scripts/ci/amd/amd_ci_start_container_disagg.sh b/scripts/ci/amd/amd_ci_start_container_disagg.sh new file mode 100755 index 000000000..ecf24f652 --- /dev/null +++ b/scripts/ci/amd/amd_ci_start_container_disagg.sh @@ -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 don’t 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}" diff --git a/test/registered/amd/disaggregation/test_disaggregation_basic.py b/test/registered/amd/disaggregation/test_disaggregation_basic.py new file mode 100644 index 000000000..be2a95c58 --- /dev/null +++ b/test/registered/amd/disaggregation/test_disaggregation_basic.py @@ -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() diff --git a/test/registered/amd/disaggregation/test_disaggregation_pp.py b/test/registered/amd/disaggregation/test_disaggregation_pp.py new file mode 100644 index 000000000..9c1dfa815 --- /dev/null +++ b/test/registered/amd/disaggregation/test_disaggregation_pp.py @@ -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() diff --git a/test/run_suite.py b/test/run_suite.py index 95e0d93bb..57b0ccff1 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -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",