run_suite.py 16.3 KB
Newer Older
1
2
import argparse
import glob
Lianmin Zheng's avatar
Lianmin Zheng committed
3
from dataclasses import dataclass
4
5
6

from sglang.test.test_utils import run_unittest_files

Lianmin Zheng's avatar
Lianmin Zheng committed
7
8
9
10
11
12
13

@dataclass
class TestFile:
    name: str
    estimated_time: float = 60


14
suites = {
15
    "per-commit": [
16
17
18
19
20
21
22
23
        TestFile("hicache/test_hicache.py", 116),
        TestFile("hicache/test_hicache_mla.py", 127),
        TestFile("hicache/test_hicache_storage.py", 127),
        TestFile("lora/test_lora.py", 200),
        TestFile("lora/test_lora_eviction.py", 200),
        TestFile("lora/test_lora_backend.py", 99),
        TestFile("lora/test_multi_lora_backend.py", 60),
        TestFile("lora/test_lora_cuda_graph.py", 250),
24
        TestFile("lora/test_lora_update.py", 400),
25
        TestFile("lora/test_lora_qwen3.py", 97),
26
        TestFile("lora/test_lora_radix_cache.py", 100),
27
        TestFile("models/test_embedding_models.py", 73),
28
        # TestFile("models/test_clip_models.py", 52),
woodx's avatar
woodx committed
29
30
        TestFile("models/test_encoder_embedding_models.py", 100),
        TestFile("models/test_cross_encoder_models.py", 100),
Lianmin Zheng's avatar
Lianmin Zheng committed
31
        TestFile("models/test_compressed_tensors_models.py", 42),
Lianmin Zheng's avatar
Lianmin Zheng committed
32
        TestFile("models/test_generation_models.py", 103),
33
        # TestFile("models/test_gme_qwen_models.py", 45),
Lianmin Zheng's avatar
Lianmin Zheng committed
34
        # TestFile("models/test_grok_models.py", 60),  # Disabled due to illegal memory access
Lianmin Zheng's avatar
Lianmin Zheng committed
35
        TestFile("models/test_qwen_models.py", 82),
Lianmin Zheng's avatar
Lianmin Zheng committed
36
        TestFile("models/test_reward_models.py", 132),
37
        TestFile("models/test_vlm_models.py", 437),
38
        TestFile("models/test_transformers_models.py", 320),
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
        TestFile("openai_server/basic/test_protocol.py", 10),
        TestFile("openai_server/basic/test_serving_chat.py", 10),
        TestFile("openai_server/basic/test_serving_completions.py", 10),
        TestFile("openai_server/basic/test_serving_embedding.py", 10),
        TestFile("openai_server/basic/test_openai_embedding.py", 141),
        TestFile("openai_server/basic/test_openai_server.py", 149),
        TestFile("openai_server/features/test_enable_thinking.py", 70),
        TestFile("openai_server/features/test_json_constrained.py", 98),
        TestFile("openai_server/features/test_json_mode.py", 90),
        TestFile("openai_server/features/test_openai_server_ebnf.py", 95),
        TestFile("openai_server/features/test_openai_server_hidden_states.py", 240),
        TestFile("openai_server/features/test_reasoning_content.py", 89),
        TestFile("openai_server/function_call/test_openai_function_calling.py", 60),
        TestFile("openai_server/function_call/test_tool_choice.py", 226),
        TestFile("openai_server/validation/test_large_max_new_tokens.py", 41),
        TestFile("openai_server/validation/test_matched_stop.py", 60),
        TestFile("openai_server/validation/test_openai_server_ignore_eos.py", 85),
        TestFile("openai_server/validation/test_request_length_validation.py", 31),
57
58
59
        TestFile("quant/test_block_int8.py", 22),
        TestFile("quant/test_fp8_kernel.py", 8),
        TestFile("quant/test_int8_kernel.py", 8),
60
        TestFile("quant/test_triton_scaled_mm.py", 8),
61
62
63
        TestFile("quant/test_w8a8_quantization.py", 46),
        TestFile("rl/test_update_weights_from_disk.py", 114),
        TestFile("rl/test_update_weights_from_tensor.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
64
        TestFile("test_abort.py", 51),
Lianmin Zheng's avatar
Lianmin Zheng committed
65
        TestFile("test_create_kvindices.py", 2),
66
        TestFile("test_chunked_prefill.py", 313),
67
        TestFile("test_eagle_infer_a.py", 370),
68
        TestFile("test_eagle_infer_b.py", 700),
Lianmin Zheng's avatar
Lianmin Zheng committed
69
70
        TestFile("test_ebnf_constrained.py", 108),
        TestFile("test_eval_fp8_accuracy.py", 303),
71
        TestFile("test_fa3.py", 376),
72
        # TestFile("test_flashmla.py", 352),
73
        TestFile("test_function_call_parser.py", 10),
Lianmin Zheng's avatar
Lianmin Zheng committed
74
        TestFile("test_fused_moe.py", 30),
75
        TestFile("test_gpt_oss_1gpu.py", 600),
76
        TestFile("test_harmony_parser.py", 20),
77
        TestFile("test_hidden_states.py", 55),
78
        TestFile("test_hybrid_attn_backend.py", 100),
Lianmin Zheng's avatar
Lianmin Zheng committed
79
        TestFile("test_input_embeddings.py", 38),
80
        TestFile("test_io_struct.py", 8),
81
        TestFile("test_jinja_template_utils.py", 1),
82
        TestFile("test_metrics.py", 32),
83
        TestFile("test_mla.py", 167),
84
        TestFile("test_mla_deepseek_v3.py", 700),
85
86
87
        TestFile("test_mla_int8_deepseek_v3.py", 429),
        TestFile("test_mla_flashinfer.py", 302),
        TestFile("test_mla_fp8.py", 93),
Lianmin Zheng's avatar
Lianmin Zheng committed
88
        TestFile("test_no_chunked_prefill.py", 108),
89
        TestFile("test_no_overlap_scheduler.py", 234),
90
        TestFile("test_original_logprobs.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
91
        TestFile("test_penalty.py", 41),
Lianmin Zheng's avatar
Lianmin Zheng committed
92
        TestFile("test_page_size.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
93
        TestFile("test_pytorch_sampling_backend.py", 66),
94
        TestFile("test_radix_attention.py", 105),
Lianmin Zheng's avatar
Lianmin Zheng committed
95
        TestFile("test_regex_constrained.py", 64),
96
        TestFile("test_reasoning_parser.py", 5),
Lianmin Zheng's avatar
Lianmin Zheng committed
97
        TestFile("test_retract_decode.py", 54),
98
        TestFile("test_request_queue_validation.py", 30),
Lianmin Zheng's avatar
Lianmin Zheng committed
99
        TestFile("test_server_args.py", 1),
100
        TestFile("test_skip_tokenizer_init.py", 117),
101
        TestFile("test_srt_engine.py", 261),
Lianmin Zheng's avatar
Lianmin Zheng committed
102
        TestFile("test_srt_endpoint.py", 130),
103
        TestFile("test_start_profile.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
104
        TestFile("test_torch_compile.py", 76),
Lianmin Zheng's avatar
Lianmin Zheng committed
105
        TestFile("test_torch_compile_moe.py", 172),
106
        TestFile("test_torch_native_attention_backend.py", 123),
Lianmin Zheng's avatar
Lianmin Zheng committed
107
108
        TestFile("test_torchao.py", 70),
        TestFile("test_triton_attention_kernels.py", 4),
109
        TestFile("test_triton_attention_backend.py", 150),
Lianmin Zheng's avatar
Lianmin Zheng committed
110
        TestFile("test_triton_moe_channel_fp8_kernel.py", 25),
111
        TestFile("test_triton_sliding_window.py", 250),
112
        TestFile("test_utils_update_weights.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
113
        TestFile("test_vision_chunked_prefill.py", 175),
114
        TestFile("test_vlm_input_format.py", 300),
115
116
        TestFile("test_vision_openai_server_a.py", 403),
        TestFile("test_vision_openai_server_b.py", 446),
117
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
118
    "per-commit-2-gpu": [
119
120
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
121
122
123
124
125
126
127
128
129
130
131
132
133
        TestFile("test_data_parallelism.py", 73),
        TestFile("test_dp_attention.py", 277),
        TestFile("test_patch_torch.py", 19),
        TestFile("test_release_memory_occupation.py", 127),
    ],
    "per-commit-4-gpu": [
        TestFile("test_gpt_oss_4gpu.py", 600),
        TestFile("test_local_attn.py", 250),
        TestFile("test_pp_single_node.py", 372),
        TestFile("test_multi_instance_release_memory_occupation.py", 64),
    ],
    "per-commit-8-gpu": [
        # Disabled because it hangs on the CI.
134
        # TestFile("ep/test_moe_ep.py", 181),
Lianmin Zheng's avatar
Lianmin Zheng committed
135
136
137
138
139
140
141
142
        TestFile("test_disaggregation.py", 499),
        TestFile("test_disaggregation_different_tp.py", 155),
        TestFile("test_full_deepseek_v3.py", 333),
    ],
    "per-commit-8-gpu-b200": [
        # add more here
    ],
    "per-commit-4-gpu-deepep": [
143
        TestFile("ep/test_deepep_small.py", 531),
Lianmin Zheng's avatar
Lianmin Zheng committed
144
145
    ],
    "per-commit-8-gpu-deepep": [
146
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
147
    ],
148
149
150
    "per-commit-8-gpu-h20": [
        TestFile("quant/test_w4a8_deepseek_v3.py", 371),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
151
152
153
154
    "nightly": [
        TestFile("test_nightly_gsm8k_eval.py"),
    ],
    "vllm_dependency_test": [
155
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
156
157
158
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
159
        # TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
160
161
162
163
164
    ],
}

# Add AMD tests
suite_amd = {
165
    "per-commit-amd": [
166
167
168
        TestFile("hicache/test_hicache.py", 116),
        TestFile("hicache/test_hicache_mla.py", 127),
        TestFile("hicache/test_hicache_storage.py", 127),
169
170
        TestFile("lora/test_lora.py", 200),
        TestFile("lora/test_lora_eviction.py", 200),
171
172
173
        TestFile("lora/test_lora_backend.py", 99),
        TestFile("lora/test_multi_lora_backend.py", 60),
        TestFile("lora/test_lora_cuda_graph.py", 250),
174
175
176
        TestFile("lora/test_lora_qwen3.py", 97),
        TestFile("models/test_embedding_models.py", 73),
        TestFile("models/test_compressed_tensors_models.py", 42),
177
178
        TestFile("models/test_qwen_models.py", 82),
        TestFile("models/test_reward_models.py", 132),
179
180
181
182
183
        TestFile("models/test_transformers_models.py", 320),
        TestFile("openai_server/basic/test_protocol.py", 10),
        TestFile("openai_server/basic/test_serving_chat.py", 10),
        TestFile("openai_server/basic/test_serving_completions.py", 10),
        TestFile("openai_server/basic/test_serving_embedding.py", 10),
184
        TestFile("openai_server/basic/test_openai_embedding.py", 141),
185
        TestFile("openai_server/basic/test_openai_server.py", 149),
186
        TestFile("openai_server/features/test_enable_thinking.py", 70),
187
188
189
190
        TestFile("openai_server/features/test_json_constrained.py", 98),
        TestFile("openai_server/features/test_json_mode.py", 90),
        TestFile("openai_server/features/test_openai_server_ebnf.py", 95),
        # TestFile("openai_server/features/test_openai_server_hidden_states.py", 240),
191
        TestFile("openai_server/features/test_reasoning_content.py", 89),
192
193
        TestFile("openai_server/function_call/test_openai_function_calling.py", 60),
        TestFile("openai_server/function_call/test_tool_choice.py", 226),
194
        TestFile("openai_server/validation/test_large_max_new_tokens.py", 41),
195
196
        TestFile("openai_server/validation/test_matched_stop.py", 60),
        TestFile("openai_server/validation/test_openai_server_ignore_eos.py", 85),
197
        TestFile("openai_server/validation/test_request_length_validation.py", 31),
198
199
200
        TestFile("quant/test_block_int8.py", 22),
        TestFile("quant/test_awq_dequant.py", 2),
        TestFile("rl/test_update_weights_from_disk.py", 114),
201
        # TestFile("rl/test_update_weights_from_tensor.py", 48),
202
203
204
        TestFile("test_abort.py", 51),
        TestFile("test_create_kvindices.py", 2),
        TestFile("test_chunked_prefill.py", 313),
205
        TestFile("test_ebnf_constrained.py", 108),
206
207
        TestFile("test_eval_fp8_accuracy.py", 303),
        TestFile("test_function_call_parser.py", 10),
208
        TestFile("test_fused_moe.py", 30),
209
        TestFile("test_input_embeddings.py", 38),
210
211
212
        TestFile("test_io_struct.py", 8),
        TestFile("test_jinja_template_utils.py", 1),
        TestFile("test_metrics.py", 32),
213
214
        TestFile("test_mla.py", 242),
        TestFile("test_mla_deepseek_v3.py", 221),
215
        TestFile("test_no_chunked_prefill.py", 108),
216
        # TestFile("test_no_overlap_scheduler.py", 234), # Disabled temporarily and track in #7703
217
218
219
220
        TestFile("test_penalty.py", 41),
        TestFile("test_page_size.py", 60),
        TestFile("test_pytorch_sampling_backend.py", 66),
        TestFile("test_radix_attention.py", 105),
221
        TestFile("test_regex_constrained.py", 64),
222
        TestFile("test_retract_decode.py", 54),
223
224
        TestFile("test_reasoning_parser.py", 5),
        TestFile("test_rope_rocm.py", 3),
225
226
        TestFile("test_server_args.py", 1),
        TestFile("test_skip_tokenizer_init.py", 117),
227
228
        TestFile("test_srt_engine.py", 261),
        TestFile("test_srt_endpoint.py", 130),
229
230
        TestFile("test_torch_compile.py", 76),
        TestFile("test_torch_compile_moe.py", 172),
231
232
        TestFile("test_torch_native_attention_backend.py", 123),
        TestFile("test_triton_attention_backend.py", 150),
233
        # TestFile("test_vision_chunked_prefill.py", 175), # Disabled temporarily and track in #7701
234
235
        TestFile("test_wave_attention_kernels.py", 2),
        TestFile("test_wave_attention_backend.py", 150),
236
237
    ],
    "per-commit-2-gpu-amd": [
238
239
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
240
241
        TestFile("test_data_parallelism.py", 73),
        TestFile("test_patch_torch.py", 19),
242
    ],
243
244
245
    "per-commit-4-gpu-amd": [
        TestFile("test_pp_single_node.py", 150),
    ],
246
247
248
    "per-commit-8-gpu-amd": [
        TestFile("test_full_deepseek_v3.py", 250),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
249
250
    "nightly-amd": [
        TestFile("test_nightly_gsm8k_eval_amd.py"),
fzyzcjy's avatar
fzyzcjy committed
251
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
252
253
254
255
}

# Add Intel Xeon tests
suite_xeon = {
256
257
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
258
        TestFile("cpu/test_binding.py"),
259
260
261
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
262
        TestFile("cpu/test_mla.py"),
263
264
265
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
266
        TestFile("cpu/test_rope.py"),
267
        TestFile("cpu/test_shared_expert.py"),
268
        TestFile("cpu/test_topk.py"),
YanbingJiang's avatar
YanbingJiang committed
269
        TestFile("test_intel_amx_attention_backend.py"),
270
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
271
272
273
274
275
}

# Add Ascend NPU tests
suite_ascend = {
    "per-commit-1-ascend-npu": [
276
        TestFile("ascend/test_ascend_tp1_bf16.py", 400),
277
        TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400),
278
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
279
    "per-commit-2-ascend-npu": [
280
        TestFile("ascend/test_ascend_tp2_bf16.py", 400),
281
        TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),
282
283
        TestFile("ascend/test_ascend_tp2_fia_bf16.py", 400),
        TestFile("ascend/test_ascend_mla_fia_w8a8int8.py", 400),
284
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
285
    "per-commit-4-ascend-npu": [
286
        TestFile("ascend/test_ascend_mla_w8a8int8.py", 400),
287
    ],
288
289
}

Lianmin Zheng's avatar
Lianmin Zheng committed
290
291
292
293
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)

Lianmin Zheng's avatar
Lianmin Zheng committed
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337

def auto_partition(files, rank, size):
    """
    Partition files into size sublists with approximately equal sums of estimated times
    using stable sorting, and return the partition for the specified rank.

    Args:
        files (list): List of file objects with estimated_time attribute
        rank (int): Index of the partition to return (0 to size-1)
        size (int): Number of partitions

    Returns:
        list: List of file objects in the specified rank's partition
    """
    weights = [f.estimated_time for f in files]

    if not weights or size <= 0 or size > len(weights):
        return []

    # Create list of (weight, original_index) tuples
    # Using negative index as secondary key to maintain original order for equal weights
    indexed_weights = [(w, -i) for i, w in enumerate(weights)]
    # Stable sort in descending order by weight
    # If weights are equal, larger (negative) index comes first (i.e., earlier original position)
    indexed_weights = sorted(indexed_weights, reverse=True)

    # Extract original indices (negate back to positive)
    indexed_weights = [(w, -i) for w, i in indexed_weights]

    # Initialize partitions and their sums
    partitions = [[] for _ in range(size)]
    sums = [0.0] * size

    # Greedy approach: assign each weight to partition with smallest current sum
    for weight, idx in indexed_weights:
        # Find partition with minimum sum
        min_sum_idx = sums.index(min(sums))
        partitions[min_sum_idx].append(idx)
        sums[min_sum_idx] += weight

    # Return the files corresponding to the indices in the specified rank's partition
    indices = partitions[rank]
    return [files[i] for i in indices]

338
339
340
341
342
343

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--timeout-per-file",
        type=int,
344
        default=1800,
345
346
347
348
349
350
351
352
353
        help="The time limit for running one file in seconds.",
    )
    arg_parser.add_argument(
        "--suite",
        type=str,
        default=list(suites.keys())[0],
        choices=list(suites.keys()) + ["all"],
        help="The suite to run",
    )
354
355
356
357
358
359
360
361
362
363
364
365
    arg_parser.add_argument(
        "--range-begin",
        type=int,
        default=0,
        help="The begin index of the range of the files to run.",
    )
    arg_parser.add_argument(
        "--range-end",
        type=int,
        default=None,
        help="The end index of the range of the files to run.",
    )
Lianmin Zheng's avatar
Lianmin Zheng committed
366
367
368
369
370
371
372
373
374
375
    arg_parser.add_argument(
        "--auto-partition-id",
        type=int,
        help="Use auto load balancing. The part id.",
    )
    arg_parser.add_argument(
        "--auto-partition-size",
        type=int,
        help="Use auto load balancing. The number of parts.",
    )
376
    args = arg_parser.parse_args()
Lianmin Zheng's avatar
Lianmin Zheng committed
377
    print(f"{args=}")
378
379
380
381
382
383

    if args.suite == "all":
        files = glob.glob("**/test_*.py", recursive=True)
    else:
        files = suites[args.suite]

Lianmin Zheng's avatar
Lianmin Zheng committed
384
385
386
387
    if args.auto_partition_size:
        files = auto_partition(files, args.auto_partition_id, args.auto_partition_size)
    else:
        files = files[args.range_begin : args.range_end]
388

Lianmin Zheng's avatar
Lianmin Zheng committed
389
    print("The running tests are ", [f.name for f in files])
390

391
392
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)