Add nightly test support to unified run_suite.py (#13941)

.github/workflows/nightly-test-nvidia.yml

@@ -8,6 +8,7 @@ on:
       - main
     paths:
       - "python/sglang/version.py"
+      - "test/run_suite.py"
   workflow_dispatch:
   workflow_call:
     inputs:
@@ -40,7 +41,7 @@ jobs:
         timeout-minutes: 60
         run: |
           cd test
-          python3 run_suite_nightly.py --suite nightly-1-gpu --continue-on-error
+          python3 run_suite.py --hw cuda --suite nightly-1-gpu --nightly --continue-on-error
 
   # General tests - 4 GPU H100
   nightly-test-general-4-gpu-h100:
@@ -60,7 +61,7 @@ jobs:
         timeout-minutes: 30
         run: |
           cd test
-          python3 run_suite_nightly.py --suite nightly-4-gpu --continue-on-error
+          python3 run_suite.py --hw cuda --suite nightly-4-gpu --nightly --continue-on-error
 
   # General tests - 8 GPU H200
   nightly-test-general-8-gpu-h200:
@@ -84,7 +85,7 @@ jobs:
           GPU_CONFIG: "8-gpu-h200"
         run: |
           cd test
-          python3 run_suite_nightly.py --suite nightly-8-gpu-h200 --continue-on-error
+          python3 run_suite.py --hw cuda --suite nightly-8-gpu-h200 --nightly --continue-on-error
 
       - name: Run Qwen3-235B nightly performance test
         timeout-minutes: 180
@@ -184,7 +185,7 @@ jobs:
           GPU_CONFIG: "8-gpu-h20"
         run: |
           cd test
-          python3 run_suite_nightly.py --suite nightly-8-gpu-h20 --continue-on-error
+          python3 run_suite.py --hw cuda --suite nightly-8-gpu-h20 --nightly --continue-on-error
 
   # Text model accuracy tests
   nightly-test-text-accuracy-2-gpu-runner:
@@ -376,7 +377,7 @@ jobs:
         timeout-minutes: 60
         run: |
           cd test
-          python3 run_suite_nightly.py --suite nightly-4-gpu-b200 --continue-on-error
+          python3 run_suite.py --hw cuda --suite nightly-4-gpu-b200 --nightly --continue-on-error
 
   # B200 Performance tests - 8 GPU
   nightly-test-perf-8-gpu-b200:

test/nightly/test_batch_invariant_ops.py

@@ -5,6 +5,9 @@ import unittest
 import torch
 
 from sglang.srt.batch_invariant_ops import batch_invariant_ops
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, suite="nightly-1-gpu", nightly=True)
 from sglang.srt.batch_invariant_ops.batch_invariant_ops import set_batch_invariant_mode
 from sglang.test.test_utils import CustomTestCase
 

test/nightly/test_cpp_radix_cache.py

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 
 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.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST,
@@ -12,6 +13,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True)
+
 
 class TestCppRadixCache(CustomTestCase):
     @classmethod

test/nightly/test_deepseek_r1_fp8_trtllm_backend.py

@@ -3,6 +3,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 as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -12,6 +13,8 @@ from sglang.test.test_utils import (
     try_cached_model,
 )
 
+register_cuda_ci(est_time=3600, suite="nightly-8-gpu-b200", nightly=True)
+
 FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
 
 

test/nightly/test_deepseek_v32_nsabackend.py

@@ -3,6 +3,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 as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -13,6 +14,8 @@ from sglang.test.test_utils import (
     write_github_step_summary,
 )
 
+register_cuda_ci(est_time=600, suite="nightly-8-gpu-h200", nightly=True)
+
 DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
 
 # Global list to collect results

test/nightly/test_deepseek_v3_deterministic.py

@@ -6,11 +6,14 @@ python3 -m unittest test_deepseek_v3_deterministic.TestFa3Deterministic
 
 import unittest
 
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_deterministic_utils import (
     COMMON_SERVER_ARGS,
     TestDeterministicBase,
 )
 
+register_cuda_ci(est_time=240, suite="nightly-1-gpu", nightly=True)
+
 DEEPSEEK_MODEL = "lmsys/sglang-ci-dsv3-test"
 
 

test/nightly/test_deepseek_v3_fp4_cutlass_moe.py

@@ -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 as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
@@ -11,6 +12,8 @@ from sglang.test.test_utils import (
     write_github_step_summary,
 )
 
+register_cuda_ci(est_time=900, suite="nightly-4-gpu-b200", nightly=True)
+
 FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4"
 SERVER_LAUNCH_TIMEOUT = 1000
 

test/nightly/test_encoder_dp.py

@@ -9,6 +9,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.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -17,6 +18,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=500, suite="nightly-4-gpu", nightly=True)
+
 MODELS = [
     SimpleNamespace(model="Qwen/Qwen2.5-VL-72B-Instruct", mmmu_accuracy=0.55),
     SimpleNamespace(model="OpenGVLab/InternVL2_5-8B", mmmu_accuracy=0.52),

test/nightly/test_flashinfer_trtllm_gen_attn_backend.py

@@ -3,6 +3,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,
@@ -11,6 +12,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=300, suite="nightly-4-gpu-b200", nightly=True)
+
 
 class TestFlashinferTrtllmGenAttnBackend(CustomTestCase):
     @classmethod

test/nightly/test_flashinfer_trtllm_gen_moe_backend.py

@@ -3,6 +3,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,
@@ -11,6 +12,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=300, suite="nightly-4-gpu-b200", nightly=True)
+
 
 class TestFlashinferTrtllmGenMoeBackend(CustomTestCase):
     @classmethod

test/nightly/test_fp4_moe.py

@@ -4,6 +4,10 @@ from typing import Callable
 import pytest
 import torch
 from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=300, suite="nightly-4-gpu-b200", nightly=True)
 from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
 from sgl_kernel import scaled_fp4_quant, silu_and_mul
 from torch.nn import functional as F

test/nightly/test_gpt_oss_4gpu_perf.py

@@ -2,8 +2,11 @@ import unittest
 
 from nightly_utils import NightlyBenchmarkRunner
 
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import DEFAULT_URL_FOR_TEST
 
+register_cuda_ci(est_time=600, suite="nightly-4-gpu-b200", nightly=True)
+
 PROFILE_DIR = "performance_profiles_gpt_oss_4gpu"
 
 

test/nightly/test_lora_eviction_policy.py

@@ -20,6 +20,9 @@ Tests LRU and FIFO eviction behavior.
 import unittest
 
 from sglang.srt.lora.eviction_policy import get_eviction_policy
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=200, suite="nightly-1-gpu", nightly=True)
 
 
 class TestLoRAEvictionPolicy(unittest.TestCase):

test/nightly/test_lora_openai_api.py

@@ -9,6 +9,9 @@ import unittest
 from unittest.mock import MagicMock
 
 from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="nightly-1-gpu", nightly=True)
 from sglang.srt.server_args import ServerArgs
 
 

test/nightly/test_lora_openai_compatible.py

@@ -17,6 +17,10 @@ import unittest
 
 import openai
 
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=150, suite="nightly-1-gpu", nightly=True)
+
 from sglang.srt.utils import kill_process_tree
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,

test/nightly/test_lora_qwen3.py

@@ -17,6 +17,10 @@ import unittest
 
 from utils import LoRAAdaptor, LoRAModelCase, run_lora_multiple_batch_on_model_cases
 
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=97, suite="nightly-1-gpu", nightly=True)
+
 from sglang.test.test_utils import CustomTestCase
 
 LORA_MODELS_QWEN3 = [

test/nightly/test_lora_radix_cache.py

@@ -18,6 +18,10 @@ import unittest
 import torch
 from utils import CI_MULTI_LORA_MODELS, run_lora_test_one_by_one
 
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=200, suite="nightly-1-gpu", nightly=True)
+
 from sglang.test.test_utils import CustomTestCase
 
 PROMPTS = [

test/nightly/test_nsa_indexer.py

@@ -4,6 +4,10 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=2, suite="nightly-1-gpu", nightly=True)
+
 from sglang.srt.layers import dp_attention as _dp_attn
 
 # Patch DP-attention globals before importing backends

test/nightly/test_qwen3_next_deterministic.py

@@ -6,11 +6,14 @@ python3 -m unittest test_qwen3_next_deterministic.TestFlashInferDeterministic
 
 import unittest
 
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_deterministic_utils import (
     COMMON_SERVER_ARGS,
     TestDeterministicBase,
 )
 
+register_cuda_ci(est_time=200, suite="nightly-4-gpu", nightly=True)
+
 QWEN3_NEXT = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 
 

test/run_suite.py

@@ -1,5 +1,6 @@
 import argparse
 import glob
+import sys
 from typing import List
 
 from sglang.test.ci.ci_register import CIRegistry, HWBackend, collect_tests
@@ -11,12 +12,29 @@ HW_MAPPING = {
     "amd": HWBackend.AMD,
 }
 
+# Per-commit test suites (run on every PR)
 PER_COMMIT_SUITES = {
     HWBackend.CPU: ["default"],
     HWBackend.AMD: ["stage-a-test-1"],
     HWBackend.CUDA: ["stage-a-test-1"],
 }
 
+# Nightly test suites (run nightly, organized by GPU configuration)
+NIGHTLY_SUITES = {
+    HWBackend.CUDA: [
+        "nightly-1-gpu",
+        "nightly-2-gpu",
+        "nightly-4-gpu",
+        "nightly-4-gpu-b200",
+        "nightly-8-gpu",
+        "nightly-8-gpu-h200",
+        "nightly-8-gpu-h20",
+        "nightly-8-gpu-b200",
+    ],
+    HWBackend.AMD: ["nightly-amd"],
+    HWBackend.CPU: [],
+}
+
 
 def filter_tests(
     ci_tests: List[CIRegistry], hw: HWBackend, suite: str, nightly: bool = False
@@ -28,14 +46,16 @@ def filter_tests(
     ]
 
     ret = []
-    for t in ci_tests:
-        if not nightly:
-            assert (
-                t.suite in PER_COMMIT_SUITES[hw]
-            ), f"Unknown stage {t.suite} for backend {hw}"
-        else:
-            raise NotImplementedError("Nightly tests are not implemented yet.")
+    valid_suites = (
+        NIGHTLY_SUITES.get(hw, []) if nightly else PER_COMMIT_SUITES.get(hw, [])
+    )
 
+    if suite not in valid_suites:
+        print(
+            f"Warning: Unknown suite {suite} for backend {hw.name}, nightly={nightly}"
+        )
+
+    for t in ci_tests:
         if t.disabled is None:
             ret.append(t)
             print(f"Including test {t.filename}")
@@ -45,20 +65,67 @@ def filter_tests(
     return ret
 
 
-def run_a_suite(hw: HWBackend, suite: str, nightly: bool = False):
-    files = glob.glob("registered/**/*.py", recursive=True)
-    ci_tests = filter_tests(collect_tests(files), hw, suite, nightly)
+def auto_partition(files, rank, size):
+    """
+    Partition files into size sublists with approximately equal sums of estimated times
+    using a greedy algorithm (LPT heuristic), and return the partition for the specified rank.
+    """
+    if not files or size <= 0:
+        return []
+
+    # Sort files by estimated_time in descending order (LPT heuristic)
+    sorted_files = sorted(files, key=lambda f: f.estimated_time, reverse=True)
+
+    partitions = [[] for _ in range(size)]
+    partition_sums = [0.0] * size
+
+    # Greedily assign each file to the partition with the smallest current total time
+    for file in sorted_files:
+        min_sum_idx = min(range(size), key=partition_sums.__getitem__)
+        partitions[min_sum_idx].append(file)
+        partition_sums[min_sum_idx] += file.estimated_time
+
+    if rank < size:
+        return partitions[rank]
+    return []
+
+
+def run_a_suite(args):
+    hw = HW_MAPPING[args.hw]
+    suite = args.suite
+    nightly = args.nightly
+    auto_partition_id = args.auto_partition_id
+    auto_partition_size = args.auto_partition_size
+
+    files = glob.glob("**/*.py", recursive=True)
+    ci_tests = filter_tests(
+        collect_tests(files, sanity_check=False), hw, suite, nightly
+    )
     test_files = [TestFile(t.filename, t.est_time) for t in ci_tests]
 
-    run_unittest_files(
+    if not test_files:
+        print(f"No tests found for hw={hw.name}, suite={suite}, nightly={nightly}")
+        print("This is expected during incremental migration. Skipping.")
+        return 0
+
+    if auto_partition_size:
+        test_files = auto_partition(test_files, auto_partition_id, auto_partition_size)
+
+    print(
+        f"Running {len(test_files)} test(s) for hw={hw.name}, suite={suite}, nightly={nightly}"
+    )
+
+    return run_unittest_files(
         test_files,
-        timeout_per_file=1200,
-        continue_on_error=False,
+        timeout_per_file=args.timeout_per_file,
+        continue_on_error=args.continue_on_error,
     )
 
 
 def main():
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(
+        description="Run CI test suites from test/registered/"
+    )
     parser.add_argument(
         "--hw",
         type=str,
@@ -67,10 +134,51 @@ def main():
         help="Hardware backend to run tests on.",
     )
     parser.add_argument("--suite", type=str, required=True, help="Test suite to run.")
-    parser.add_argument("--nightly", action="store_true")
+    parser.add_argument(
+        "--nightly",
+        action="store_true",
+        help="Run nightly tests instead of per-commit tests.",
+    )
+    parser.add_argument(
+        "--timeout-per-file",
+        type=int,
+        default=1200,
+        help="The time limit for running one file in seconds (default: 1200).",
+    )
+    parser.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        default=False,
+        help="Continue running remaining tests even if one fails (default: False, useful for nightly tests).",
+    )
+    parser.add_argument(
+        "--auto-partition-id",
+        type=int,
+        help="Use auto load balancing. The part id.",
+    )
+    parser.add_argument(
+        "--auto-partition-size",
+        type=int,
+        help="Use auto load balancing. The number of parts.",
+    )
     args = parser.parse_args()
-    hw = HW_MAPPING[args.hw]
-    run_a_suite(hw, args.suite, args.nightly)
+
+    # Validate auto-partition arguments
+    if (args.auto_partition_id is not None) != (args.auto_partition_size is not None):
+        parser.error(
+            "--auto-partition-id and --auto-partition-size must be specified together."
+        )
+    if args.auto_partition_size is not None:
+        if args.auto_partition_size <= 0:
+            parser.error("--auto-partition-size must be positive.")
+        if not 0 <= args.auto_partition_id < args.auto_partition_size:
+            parser.error(
+                f"--auto-partition-id must be in range [0, {args.auto_partition_size}), "
+                f"but got {args.auto_partition_id}"
+            )
+
+    exit_code = run_a_suite(args)
+    sys.exit(exit_code)
 
 
 if __name__ == "__main__":