run_suite.py 17.6 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
        TestFile("function_call/test_json_schema_constraint.py", 30),
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),
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", 741),
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),
79
        TestFile("test_standalone_speculative_decoding.py", 250),
80
        TestFile("test_ngram_speculative_decoding.py", 250),
Lianmin Zheng's avatar
Lianmin Zheng committed
81
        TestFile("test_input_embeddings.py", 38),
82
        TestFile("test_io_struct.py", 8),
83
        TestFile("test_jinja_template_utils.py", 1),
84
        TestFile("test_logprobs.py", 55),
85
        TestFile("test_metrics.py", 32),
86
        TestFile("test_metrics_utils.py", 1),
87
        TestFile("test_mla.py", 167),
Shangming Cai's avatar
Shangming Cai committed
88
        TestFile("test_mla_deepseek_v3.py", 1420),
89
90
91
        TestFile("test_mla_int8_deepseek_v3.py", 429),
        TestFile("test_mla_flashinfer.py", 302),
        TestFile("test_mla_fp8.py", 93),
92
        TestFile("test_multi_tokenizer.py", 230),
Lianmin Zheng's avatar
Lianmin Zheng committed
93
        TestFile("test_no_chunked_prefill.py", 108),
94
        TestFile("test_no_overlap_scheduler.py", 234),
95
        TestFile("test_original_logprobs.py", 41),
Lianmin Zheng's avatar
Lianmin Zheng committed
96
        TestFile("test_penalty.py", 41),
Lianmin Zheng's avatar
Lianmin Zheng committed
97
        TestFile("test_page_size.py", 60),
98
        TestFile("test_priority_scheduling.py", 100),
Lianmin Zheng's avatar
Lianmin Zheng committed
99
        TestFile("test_pytorch_sampling_backend.py", 66),
100
        TestFile("test_radix_attention.py", 105),
101
        TestFile("test_radix_cache_unit.py", 5),
Lianmin Zheng's avatar
Lianmin Zheng committed
102
        TestFile("test_regex_constrained.py", 64),
103
        TestFile("test_reasoning_parser.py", 5),
Lianmin Zheng's avatar
Lianmin Zheng committed
104
        TestFile("test_retract_decode.py", 54),
105
        TestFile("test_request_queue_validation.py", 30),
106
        TestFile("test_score_api.py", 180),
Lianmin Zheng's avatar
Lianmin Zheng committed
107
        TestFile("test_server_args.py", 1),
108
        TestFile("test_skip_tokenizer_init.py", 117),
109
        TestFile("test_srt_engine.py", 261),
Lianmin Zheng's avatar
Lianmin Zheng committed
110
        TestFile("test_srt_endpoint.py", 130),
111
        TestFile("test_start_profile.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
112
        TestFile("test_torch_compile.py", 76),
Lianmin Zheng's avatar
Lianmin Zheng committed
113
        TestFile("test_torch_compile_moe.py", 172),
114
        TestFile("test_torch_native_attention_backend.py", 123),
Lianmin Zheng's avatar
Lianmin Zheng committed
115
116
        TestFile("test_torchao.py", 70),
        TestFile("test_triton_attention_kernels.py", 4),
117
        TestFile("test_triton_attention_backend.py", 150),
Lianmin Zheng's avatar
Lianmin Zheng committed
118
        TestFile("test_triton_moe_channel_fp8_kernel.py", 25),
119
        TestFile("test_triton_sliding_window.py", 250),
120
        TestFile("test_utils_update_weights.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
121
        TestFile("test_vision_chunked_prefill.py", 175),
122
        TestFile("test_vlm_input_format.py", 300),
Shangming Cai's avatar
Shangming Cai committed
123
        TestFile("test_vision_openai_server_a.py", 724),
124
        TestFile("test_vision_openai_server_b.py", 446),
125
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
126
    "per-commit-2-gpu": [
127
128
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
129
130
        TestFile("test_data_parallelism.py", 73),
        TestFile("test_dp_attention.py", 277),
131
        TestFile("test_load_weights_from_remote_instance.py", 72),
Lianmin Zheng's avatar
Lianmin Zheng committed
132
        TestFile("test_patch_torch.py", 19),
Shangming Cai's avatar
Shangming Cai committed
133
        TestFile("test_release_memory_occupation.py", 257),
134
135
        TestFile("hicache/test_hicache_storage_file_backend.py", 400),
        TestFile("hicache/test_hicache_storage_3fs_backend.py", 400),
Lianmin Zheng's avatar
Lianmin Zheng committed
136
137
138
139
140
    ],
    "per-commit-4-gpu": [
        TestFile("test_gpt_oss_4gpu.py", 600),
        TestFile("test_local_attn.py", 250),
        TestFile("test_pp_single_node.py", 372),
Yi Zhang's avatar
Yi Zhang committed
141
        TestFile("models/test_qwen3_next_models.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
142
143
144
        TestFile("test_multi_instance_release_memory_occupation.py", 64),
    ],
    "per-commit-8-gpu": [
145
        TestFile("hicache/test_hicache_storage_mooncake_backend.py", 800),
146
        TestFile("lora/test_lora_llama4.py", 600),
Lianmin Zheng's avatar
Lianmin Zheng committed
147
148
        TestFile("test_disaggregation.py", 499),
        TestFile("test_disaggregation_different_tp.py", 155),
149
        TestFile("test_disaggregation_pp.py", 60),
Lianmin Zheng's avatar
Lianmin Zheng committed
150
151
        TestFile("test_full_deepseek_v3.py", 333),
    ],
Yineng Zhang's avatar
Yineng Zhang committed
152
    "per-commit-4-gpu-b200": [
153
        # TestFile("test_gpt_oss_4gpu.py", 600),
154
        # TestFile("test_deepseek_v3_fp4_4gpu.py", 3600),
Lianmin Zheng's avatar
Lianmin Zheng committed
155
156
    ],
    "per-commit-4-gpu-deepep": [
157
        TestFile("ep/test_deepep_small.py", 531),
Lianmin Zheng's avatar
Lianmin Zheng committed
158
159
    ],
    "per-commit-8-gpu-deepep": [
160
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
161
    ],
162
163
164
    "per-commit-8-gpu-h20": [
        TestFile("quant/test_w4a8_deepseek_v3.py", 371),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
165
    "vllm_dependency_test": [
166
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
167
168
169
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
170
        # TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
171
172
173
174
175
    ],
}

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

# Add Intel Xeon tests
suite_xeon = {
274
275
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
276
        TestFile("cpu/test_binding.py"),
277
278
279
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
280
        TestFile("cpu/test_mla.py"),
281
282
283
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
284
        TestFile("cpu/test_rope.py"),
285
        TestFile("cpu/test_shared_expert.py"),
286
        TestFile("cpu/test_topk.py"),
YanbingJiang's avatar
YanbingJiang committed
287
        TestFile("test_intel_amx_attention_backend.py"),
288
        TestFile("test_cpu_graph.py"),
289
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
290
291
292
293
294
}

# Add Ascend NPU tests
suite_ascend = {
    "per-commit-1-ascend-npu": [
295
        TestFile("ascend/test_ascend_tp1_bf16.py", 400),
296
        TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400),
297
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
298
    "per-commit-2-ascend-npu": [
299
        TestFile("ascend/test_ascend_tp2_bf16.py", 400),
300
        TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),
301
302
        TestFile("ascend/test_ascend_tp2_fia_bf16.py", 400),
        TestFile("ascend/test_ascend_mla_fia_w8a8int8.py", 400),
303
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
304
    "per-commit-4-ascend-npu": [
305
        TestFile("ascend/test_ascend_mla_w8a8int8.py", 400),
306
        TestFile("ascend/test_ascend_tp4_bf16.py", 400),
307
    ],
308
309
310
    "per-commit-16-ascend-a3": [
        TestFile("ascend/test_ascend_deepep.py", 400),
    ],
311
312
}

Lianmin Zheng's avatar
Lianmin Zheng committed
313
314
315
316
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)

Lianmin Zheng's avatar
Lianmin Zheng committed
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
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360

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]

361
362
363
364
365
366

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--timeout-per-file",
        type=int,
367
        default=1200,
368
369
370
371
372
373
374
375
376
        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",
    )
377
378
379
380
381
382
383
384
385
386
387
388
    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
389
390
391
392
393
394
395
396
397
398
    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.",
    )
399
    args = arg_parser.parse_args()
Lianmin Zheng's avatar
Lianmin Zheng committed
400
    print(f"{args=}")
401
402
403
404
405
406

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

Lianmin Zheng's avatar
Lianmin Zheng committed
407
408
409
410
    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]
411

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

414
415
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)