run_suite.py 16.7 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_metrics_utils.py", 1),
84
        TestFile("test_mla.py", 167),
85
        TestFile("test_mla_deepseek_v3.py", 700),
86
87
88
        TestFile("test_mla_int8_deepseek_v3.py", 429),
        TestFile("test_mla_flashinfer.py", 302),
        TestFile("test_mla_fp8.py", 93),
89
        TestFile("test_multi_tokenizer.py", 230),
Lianmin Zheng's avatar
Lianmin Zheng committed
90
        TestFile("test_no_chunked_prefill.py", 108),
91
        TestFile("test_no_overlap_scheduler.py", 234),
92
        TestFile("test_original_logprobs.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
93
        TestFile("test_penalty.py", 41),
Lianmin Zheng's avatar
Lianmin Zheng committed
94
        TestFile("test_page_size.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
95
        TestFile("test_pytorch_sampling_backend.py", 66),
96
        TestFile("test_radix_attention.py", 105),
Lianmin Zheng's avatar
Lianmin Zheng committed
97
        TestFile("test_regex_constrained.py", 64),
98
        TestFile("test_reasoning_parser.py", 5),
Lianmin Zheng's avatar
Lianmin Zheng committed
99
        TestFile("test_retract_decode.py", 54),
100
        TestFile("test_request_queue_validation.py", 30),
Lianmin Zheng's avatar
Lianmin Zheng committed
101
        TestFile("test_server_args.py", 1),
102
        TestFile("test_skip_tokenizer_init.py", 117),
103
        TestFile("test_srt_engine.py", 261),
Lianmin Zheng's avatar
Lianmin Zheng committed
104
        TestFile("test_srt_endpoint.py", 130),
105
        TestFile("test_start_profile.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
106
        TestFile("test_torch_compile.py", 76),
Lianmin Zheng's avatar
Lianmin Zheng committed
107
        TestFile("test_torch_compile_moe.py", 172),
108
        TestFile("test_torch_native_attention_backend.py", 123),
Lianmin Zheng's avatar
Lianmin Zheng committed
109
110
        TestFile("test_torchao.py", 70),
        TestFile("test_triton_attention_kernels.py", 4),
111
        TestFile("test_triton_attention_backend.py", 150),
Lianmin Zheng's avatar
Lianmin Zheng committed
112
        TestFile("test_triton_moe_channel_fp8_kernel.py", 25),
113
        TestFile("test_triton_sliding_window.py", 250),
114
        TestFile("test_utils_update_weights.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
115
        TestFile("test_vision_chunked_prefill.py", 175),
116
        TestFile("test_vlm_input_format.py", 300),
117
118
        TestFile("test_vision_openai_server_a.py", 403),
        TestFile("test_vision_openai_server_b.py", 446),
119
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
120
    "per-commit-2-gpu": [
121
122
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
123
124
125
126
        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),
127
128
        TestFile("hicache/test_hicache_storage_e2e.py", 400),
        TestFile("hicache/test_hicache_storage_benchmark.py", 400),
Lianmin Zheng's avatar
Lianmin Zheng committed
129
130
131
132
133
134
135
136
137
    ],
    "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.
138
        # TestFile("ep/test_moe_ep.py", 181),
139
        TestFile("lora/test_lora_llama4.py", 600),
Lianmin Zheng's avatar
Lianmin Zheng committed
140
141
        TestFile("test_disaggregation.py", 499),
        TestFile("test_disaggregation_different_tp.py", 155),
142
        TestFile("test_disaggregation_pp.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
143
144
145
146
        TestFile("test_full_deepseek_v3.py", 333),
    ],
    "per-commit-8-gpu-b200": [
        # add more here
147
        TestFile("test_gpt_oss_4gpu.py", 600),
Lianmin Zheng's avatar
Lianmin Zheng committed
148
149
    ],
    "per-commit-4-gpu-deepep": [
150
        TestFile("ep/test_deepep_small.py", 531),
Lianmin Zheng's avatar
Lianmin Zheng committed
151
152
    ],
    "per-commit-8-gpu-deepep": [
153
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
154
    ],
155
156
157
    "per-commit-8-gpu-h20": [
        TestFile("quant/test_w4a8_deepseek_v3.py", 371),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
158
159
160
161
    "nightly": [
        TestFile("test_nightly_gsm8k_eval.py"),
    ],
    "vllm_dependency_test": [
162
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
163
164
165
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
166
        # TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
167
168
169
170
171
    ],
}

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

# Add Intel Xeon tests
suite_xeon = {
264
265
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
266
        TestFile("cpu/test_binding.py"),
267
268
269
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
270
        TestFile("cpu/test_mla.py"),
271
272
273
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
274
        TestFile("cpu/test_rope.py"),
275
        TestFile("cpu/test_shared_expert.py"),
276
        TestFile("cpu/test_topk.py"),
YanbingJiang's avatar
YanbingJiang committed
277
        TestFile("test_intel_amx_attention_backend.py"),
278
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
279
280
281
282
283
}

# Add Ascend NPU tests
suite_ascend = {
    "per-commit-1-ascend-npu": [
284
        TestFile("ascend/test_ascend_tp1_bf16.py", 400),
285
        TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400),
286
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
287
    "per-commit-2-ascend-npu": [
288
        TestFile("ascend/test_ascend_tp2_bf16.py", 400),
289
        TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),
290
291
        TestFile("ascend/test_ascend_tp2_fia_bf16.py", 400),
        TestFile("ascend/test_ascend_mla_fia_w8a8int8.py", 400),
292
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
293
    "per-commit-4-ascend-npu": [
294
        TestFile("ascend/test_ascend_mla_w8a8int8.py", 400),
295
    ],
296
297
}

Lianmin Zheng's avatar
Lianmin Zheng committed
298
299
300
301
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)

Lianmin Zheng's avatar
Lianmin Zheng committed
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
338
339
340
341
342
343
344
345

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]

346
347
348
349
350
351

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--timeout-per-file",
        type=int,
352
        default=1800,
353
354
355
356
357
358
359
360
361
        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",
    )
362
363
364
365
366
367
368
369
370
371
372
373
    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
374
375
376
377
378
379
380
381
382
383
    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.",
    )
384
    args = arg_parser.parse_args()
Lianmin Zheng's avatar
Lianmin Zheng committed
385
    print(f"{args=}")
386
387
388
389
390
391

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

Lianmin Zheng's avatar
Lianmin Zheng committed
392
393
394
395
    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]
396

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

399
400
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)