ci: migrate 2-GPU tests to test/registered/ (#16529)

This commit is contained in:
Alison Shao
2026-01-07 20:28:16 -08:00
committed by GitHub
parent ab7d5829cd
commit 63cc97f4ef
15 changed files with 84 additions and 80 deletions

View File

@@ -674,6 +674,10 @@ jobs:
runs-on: 2-gpu-runner
env:
RUNNER_LABELS: 2-gpu-runner
strategy:
fail-fast: false
matrix:
partition: [0, 1]
steps:
- name: Checkout code
uses: actions/checkout@v4
@@ -700,7 +704,7 @@ jobs:
if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
CONTINUE_ON_ERROR_FLAG="--continue-on-error"
fi
python3 run_suite.py --hw cuda --suite stage-b-test-large-2-gpu $CONTINUE_ON_ERROR_FLAG
python3 run_suite.py --hw cuda --suite stage-b-test-large-2-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 2 $CONTINUE_ON_ERROR_FLAG
stage-c-test-large-4-gpu:
needs: [check-changes, call-gate, stage-b-test-small-1-gpu, stage-b-test-large-1-gpu, stage-b-test-large-2-gpu, sgl-kernel-build-wheels]
@@ -1027,57 +1031,6 @@ jobs:
IS_BLACKWELL=1 python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 $CONTINUE_ON_ERROR_FLAG
unit-test-backend-2-gpu:
needs: [check-changes, call-gate, unit-test-backend-1-gpu]
if: |
always() &&
(
(inputs.target_stage == 'unit-test-backend-2-gpu') ||
(
!inputs.target_stage &&
(github.event_name == 'schedule' || (!failure() && !cancelled())) &&
((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
)
)
runs-on: 2-gpu-runner
env:
RUNNER_LABELS: 2-gpu-runner
strategy:
fail-fast: false
matrix:
part: [0, 1]
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
- name: Download artifacts
if: needs.check-changes.outputs.sgl_kernel == 'true'
uses: actions/download-artifact@v4
with:
path: sgl-kernel/dist/
merge-multiple: true
pattern: wheel-python3.10-cuda12.9
- name: Install dependencies
run: |
CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
- name: Run test
timeout-minutes: 30
run: |
cd test/srt
RETRY_FLAG=""
if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
RETRY_FLAG="--enable-retry"
fi
CONTINUE_ON_ERROR_FLAG=""
if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
CONTINUE_ON_ERROR_FLAG="--continue-on-error"
fi
python3 run_suite.py --suite per-commit-2-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 $RETRY_FLAG $CONTINUE_ON_ERROR_FLAG
unit-test-backend-4-gpu:
needs: [check-changes, call-gate, unit-test-backend-1-gpu, stage-b-test-4-gpu-b200]
if: |
@@ -1802,7 +1755,6 @@ jobs:
stage-c-test-large-4-gpu,
quantization-test,
unit-test-backend-1-gpu,
unit-test-backend-2-gpu,
stage-b-test-4-gpu-b200,
unit-test-backend-4-gpu,
unit-test-backend-8-gpu-h20,

View File

@@ -219,6 +219,26 @@ def safetensors_weights_iterator(
yield name, param
def _load_pt_file(bin_file: str, device: str) -> dict:
"""Load a PyTorch checkpoint file, handling legacy tar format.
PyTorch 2.6 changed the default of weights_only from False to True.
Legacy tar format files cannot be loaded with weights_only=True.
This function tries weights_only=True first, then falls back to False
for legacy tar format files from trusted sources (HuggingFace Hub).
"""
try:
return torch.load(bin_file, map_location=device, weights_only=True)
except RuntimeError as e:
if "legacy .tar format" in str(e):
logger.warning(
"Loading %s with weights_only=False (legacy tar format)",
os.path.basename(bin_file),
)
return torch.load(bin_file, map_location=device, weights_only=False)
raise
def pt_weights_iterator(
hf_weights_files: list[str],
to_cpu: bool = True,
@@ -234,7 +254,7 @@ def pt_weights_iterator(
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
):
state = torch.load(bin_file, map_location=device, weights_only=True)
state = _load_pt_file(bin_file, device)
yield from state.items()
del state

View File

@@ -65,7 +65,10 @@ _active_symmetric_memory_context = None
def is_symmetric_memory_enabled():
return get_global_server_args().enable_symm_mem
try:
return get_global_server_args().enable_symm_mem
except ValueError:
return False
def set_graph_pool_id(graph_pool_id):

View File

@@ -822,6 +822,26 @@ def multi_thread_safetensors_weights_iterator(
yield name, param
def _load_pt_file(bin_file: str) -> dict:
"""Load a PyTorch checkpoint file, handling legacy tar format.
PyTorch 2.6 changed the default of weights_only from False to True.
Legacy tar format files cannot be loaded with weights_only=True.
This function tries weights_only=True first, then falls back to False
for legacy tar format files from trusted sources (HuggingFace Hub).
"""
try:
return torch.load(bin_file, map_location="cpu", weights_only=True)
except RuntimeError as e:
if "legacy .tar format" in str(e):
logger.warning(
"Loading %s with weights_only=False (legacy tar format)",
os.path.basename(bin_file),
)
return torch.load(bin_file, map_location="cpu", weights_only=False)
raise
def pt_weights_iterator(
hf_weights_files: List[str],
) -> Generator[Tuple[str, torch.Tensor], None, None]:
@@ -835,7 +855,7 @@ def pt_weights_iterator(
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
):
state = torch.load(bin_file, map_location="cpu", weights_only=True)
state = _load_pt_file(bin_file)
yield from state.items()
del state
@@ -849,12 +869,9 @@ def multi_thread_pt_weights_iterator(
not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
)
def _load_file(bin_file: str):
return torch.load(bin_file, map_location="cpu", weights_only=True)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(_load_file, bin_file) for bin_file in hf_weights_files
executor.submit(_load_pt_file, bin_file) for bin_file in hf_weights_files
]
if enable_tqdm:

View File

@@ -155,7 +155,6 @@ def handle_rerun_stage(
"multimodal-gen-test-2-gpu",
"quantization-test",
"unit-test-backend-1-gpu",
"unit-test-backend-2-gpu",
"stage-b-test-4-gpu-b200",
"unit-test-backend-4-gpu",
"unit-test-backend-8-gpu-h200",

View File

@@ -7,6 +7,7 @@ import openai
import requests
from transformers import AutoTokenizer
from sglang.test.ci.ci_register import register_cuda_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,
@@ -19,6 +20,8 @@ from sglang.test.test_utils import (
popen_launch_pd_server,
)
register_cuda_ci(est_time=400, suite="stage-b-test-large-2-gpu")
class TestDisaggregationAccuracy(PDDisaggregationServerBase):
@classmethod

View File

@@ -5,6 +5,7 @@ from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
@@ -14,6 +15,9 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=73, suite="stage-b-test-large-2-gpu")
register_amd_ci(est_time=73, suite="stage-b-test-large-2-gpu-amd")
class TestDataParallelism(CustomTestCase):
@classmethod

View File

@@ -5,6 +5,7 @@ import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
from sglang.test.run_eval import run_eval
@@ -19,6 +20,8 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=350, suite="stage-b-test-large-2-gpu")
class TestDPAttentionDP2TP2(CustomTestCase):
@classmethod

View File

@@ -24,6 +24,7 @@ import torch
import torch.multiprocessing as mp
import sglang as sgl
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
@@ -37,6 +38,9 @@ from sglang.utils import terminate_process
mp.set_start_method("spawn", force=True)
register_cuda_ci(est_time=72, suite="stage-b-test-large-2-gpu")
register_amd_ci(est_time=72, suite="stage-b-test-large-2-gpu-amd")
def verify_params_close(params1, params2, error_msg):
"""Verify if two parameter arrays are close enough."""

View File

@@ -1,7 +1,7 @@
"""
Benchmark tests for HiCache Storage with 3FS backend.
Usage:
python3 -m pytest test/srt/hicache/test_hicache_storage_3fs_backend.py -v
python3 -m pytest test/registered/hicache/test_hicache_storage_3fs_backend.py -v
"""
import json
@@ -10,8 +10,11 @@ import unittest
from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=200, suite="stage-b-test-large-2-gpu")
class HiCacheStorage3FSBackendBaseMixin(HiCacheStorageBaseMixin):
"""Base mixin class with common setup and utilities"""

View File

@@ -1,7 +1,7 @@
"""
E2E tests for HiCache Storage functionality.
Usage:
python3 -m pytest test/srt/hicache/test_hicache_storage_e2e.py -v
python3 -m pytest test/registered/hicache/test_hicache_storage_file_backend.py -v
"""
import json
@@ -18,6 +18,7 @@ import requests
from sglang.bench_serving import get_tokenizer
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
@@ -29,6 +30,8 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=200, suite="stage-b-test-large-2-gpu")
class HiCacheStorageBaseMixin:
"""Base mixin class with common setup and utilities"""

View File

@@ -1,7 +1,7 @@
"""
Benchmark tests for HiCache Storage with Mooncake backend.
Usage:
python3.10 -m pytest test/srt/hicache/test_hicache_storage_mooncake_backend.py -v
python3.10 -m pytest test/registered/hicache/test_hicache_storage_mooncake_backend.py -v
"""
import os
@@ -12,6 +12,7 @@ import unittest
import requests
from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
CustomTestCase,
@@ -19,6 +20,8 @@ from sglang.test.test_utils import (
is_in_ci,
)
register_cuda_ci(est_time=300, suite="stage-b-test-large-2-gpu")
class HiCacheStorageMooncakeBackendBaseMixin(HiCacheStorageBaseMixin):
"""Base mixin class with common setup and utilities"""

View File

@@ -2,6 +2,7 @@ import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -10,6 +11,8 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=90, suite="stage-b-test-large-2-gpu")
class TestKimiLinear(CustomTestCase):
@classmethod

View File

@@ -1,9 +1,12 @@
import unittest
from sglang.srt.utils import is_blackwell
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=132, suite="stage-b-test-large-2-gpu")
class TestNvidiaNemotronNanoV2BF16(GSM8KMixin, DefaultServerBase):
model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"

View File

@@ -16,17 +16,6 @@ suites = {
TestFile("test_video_utils.py", 5),
TestFile("test_modelopt_export.py", 9),
],
"per-commit-2-gpu": [
TestFile("hicache/test_hicache_storage_3fs_backend.py", 200),
TestFile("hicache/test_hicache_storage_file_backend.py", 200),
TestFile("hicache/test_hicache_storage_mooncake_backend.py", 300),
TestFile("models/test_kimi_linear_models.py", 90),
TestFile("models/test_nvidia_nemotron_nano_v2.py", 132),
TestFile("test_data_parallelism.py", 73),
TestFile("test_disaggregation_basic.py", 400),
TestFile("test_dp_attention.py", 350),
TestFile("test_load_weights_from_remote_instance.py", 72),
],
"per-commit-4-gpu": [
TestFile("models/test_qwen3_next_models.py", 650),
TestFile("test_gpt_oss_4gpu.py", 300),
@@ -116,11 +105,6 @@ suite_amd = {
# TestFile("test_vision_chunked_prefill.py", 175), # Disabled temporarily and track in #7701
# TestFile("test_wave_attention_backend.py", 150), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
],
"per-commit-amd-mi35x": [],
"per-commit-2-gpu-amd": [
TestFile("test_data_parallelism.py", 73),
TestFile("test_load_weights_from_remote_instance.py", 72),
],
"per-commit-4-gpu-amd": [
TestFile("test_pp_single_node.py", 150),
],