run_suite.py 19 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
# NOTE: please sort the test cases alphabetically by the test file name
15
suites = {
Lianmin Zheng's avatar
Lianmin Zheng committed
16
    "per-commit-1-gpu": [
17
        TestFile("function_call/test_json_schema_constraint.py", 30),
18
        TestFile("hicache/test_hicache.py", 116),
19
        TestFile("hicache/test_hicache_eagle.py", 150),
20
21
22
23
        TestFile("hicache/test_hicache_mla.py", 127),
        TestFile("hicache/test_hicache_storage.py", 127),
        TestFile("lora/test_lora.py", 200),
        TestFile("lora/test_lora_backend.py", 99),
24
        TestFile("lora/test_lora_eviction.py", 200),
25
        TestFile("lora/test_lora_qwen3.py", 97),
26
        TestFile("lora/test_lora_radix_cache.py", 100),
27
28
29
30
        TestFile("lora/test_lora_update.py", 400),
        TestFile("lora/test_multi_lora_backend.py", 60),
        TestFile("models/test_compressed_tensors_models.py", 42),
        TestFile("models/test_cross_encoder_models.py", 100),
31
        TestFile("models/test_embedding_models.py", 73),
32
        TestFile("models/test_encoder_embedding_models.py", 460),
Lianmin Zheng's avatar
Lianmin Zheng committed
33
        TestFile("models/test_generation_models.py", 103),
34
        TestFile("models/test_nvidia_nemotron_nano_v2.py", 180),
Lianmin Zheng's avatar
Lianmin Zheng committed
35
        TestFile("models/test_qwen_models.py", 82),
Stefan He's avatar
Stefan He committed
36
        TestFile("batch_invariant/test_batch_invariant_ops.py", 10),
Lianmin Zheng's avatar
Lianmin Zheng committed
37
        TestFile("models/test_reward_models.py", 132),
38
        TestFile("models/test_transformers_models.py", 320),
39
40
41
        TestFile("models/test_vlm_models.py", 741),
        TestFile("openai_server/basic/test_openai_embedding.py", 141),
        TestFile("openai_server/basic/test_openai_server.py", 149),
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
        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/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),
58
59
60
        TestFile("quant/test_block_int8.py", 22),
        TestFile("quant/test_fp8_kernel.py", 8),
        TestFile("quant/test_int8_kernel.py", 8),
61
        TestFile("quant/test_triton_scaled_mm.py", 8),
62
        TestFile("quant/test_w8a8_quantization.py", 46),
63
        TestFile("rl/test_fp32_lm_head.py", 30),
64
65
        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
66
        TestFile("test_abort.py", 51),
67
        TestFile("test_chunked_prefill.py", 313),
68
        TestFile("test_create_kvindices.py", 2),
69
        TestFile("test_deterministic.py", 300),
70
        TestFile("test_eagle_infer_a.py", 370),
71
        TestFile("test_eagle_infer_b.py", 700),
72
        TestFile("test_eagle_infer_beta.py", 300),
Lianmin Zheng's avatar
Lianmin Zheng committed
73
74
        TestFile("test_ebnf_constrained.py", 108),
        TestFile("test_eval_fp8_accuracy.py", 303),
75
        TestFile("test_fa3.py", 376),
76
        # TestFile("test_flashmla.py", 352),
77
        TestFile("test_function_call_parser.py", 10),
Lianmin Zheng's avatar
Lianmin Zheng committed
78
        TestFile("test_fused_moe.py", 30),
79
        TestFile("test_gpt_oss_1gpu.py", 600),
80
        TestFile("test_harmony_parser.py", 20),
81
        TestFile("test_hidden_states.py", 55),
82
        TestFile("test_hybrid_attn_backend.py", 379),
Lianmin Zheng's avatar
Lianmin Zheng committed
83
        TestFile("test_input_embeddings.py", 38),
84
        TestFile("test_io_struct.py", 8),
85
        TestFile("test_jinja_template_utils.py", 1),
86
        TestFile("test_logprobs.py", 55),
87
        TestFile("test_mamba_unittest.py", 4),
88
        TestFile("test_metrics.py", 32),
89
        TestFile("test_metrics_utils.py", 1),
90
        TestFile("test_mla.py", 167),
Lianmin Zheng's avatar
Lianmin Zheng committed
91
        TestFile("test_mla_deepseek_v3.py", 500),
92
93
        TestFile("test_mla_flashinfer.py", 302),
        TestFile("test_mla_fp8.py", 93),
94
95
        TestFile("test_mla_int8_deepseek_v3.py", 429),
        TestFile("test_modelopt_loader.py", 30),
96
        TestFile("test_multi_tokenizer.py", 230),
97
        TestFile("test_ngram_speculative_decoding.py", 250),
Lianmin Zheng's avatar
Lianmin Zheng committed
98
        TestFile("test_no_chunked_prefill.py", 108),
99
        TestFile("test_no_overlap_scheduler.py", 234),
100
        TestFile("test_original_logprobs.py", 41),
Lianmin Zheng's avatar
Lianmin Zheng committed
101
        TestFile("test_page_size.py", 60),
102
        TestFile("test_penalty.py", 41),
103
        TestFile("test_priority_scheduling.py", 100),
Lianmin Zheng's avatar
Lianmin Zheng committed
104
        TestFile("test_pytorch_sampling_backend.py", 66),
105
        TestFile("test_radix_attention.py", 105),
106
        TestFile("test_radix_cache_unit.py", 5),
107
        TestFile("test_reasoning_parser.py", 5),
108
        TestFile("test_regex_constrained.py", 64),
109
        TestFile("test_request_queue_validation.py", 30),
110
        TestFile("test_retract_decode.py", 54),
111
        TestFile("test_score_api.py", 310),
Lianmin Zheng's avatar
Lianmin Zheng committed
112
        TestFile("test_server_args.py", 1),
113
        TestFile("test_skip_tokenizer_init.py", 117),
Lianmin Zheng's avatar
Lianmin Zheng committed
114
        TestFile("test_srt_endpoint.py", 130),
115
116
        TestFile("test_srt_engine.py", 261),
        TestFile("test_standalone_speculative_decoding.py", 250),
117
        TestFile("test_start_profile.py", 60),
118
119
        TestFile("test_profile_merger.py", 60),
        TestFile("test_profile_merger_http_api.py", 15),
120
        TestFile("test_swa_unittest.py", 1),
Lianmin Zheng's avatar
Lianmin Zheng committed
121
        TestFile("test_torch_compile.py", 76),
Lianmin Zheng's avatar
Lianmin Zheng committed
122
        TestFile("test_torch_compile_moe.py", 172),
123
        TestFile("test_torch_native_attention_backend.py", 123),
Lianmin Zheng's avatar
Lianmin Zheng committed
124
        TestFile("test_torchao.py", 70),
125
        TestFile("test_triton_attention_backend.py", 150),
126
        TestFile("test_triton_attention_kernels.py", 4),
Lianmin Zheng's avatar
Lianmin Zheng committed
127
        TestFile("test_triton_moe_channel_fp8_kernel.py", 25),
128
        TestFile("test_triton_sliding_window.py", 250),
129
        TestFile("test_utils_update_weights.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
130
        TestFile("test_vision_chunked_prefill.py", 175),
Shangming Cai's avatar
Shangming Cai committed
131
        TestFile("test_vision_openai_server_a.py", 724),
132
        TestFile("test_vision_openai_server_b.py", 446),
133
        TestFile("test_vlm_input_format.py", 300),
134
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
135
    "per-commit-2-gpu": [
136
        TestFile("ep/test_moe_ep.py", 140),
137
138
139
140
        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", 400),
        TestFile("layers/attention/mamba/test_mamba2_mixer.py", 110),
141
142
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
143
        TestFile("test_data_parallelism.py", 73),
144
        TestFile("test_disaggregation_basic.py", 400),
Lianmin Zheng's avatar
Lianmin Zheng committed
145
        TestFile("test_dp_attention.py", 594),
146
        TestFile("test_load_weights_from_remote_instance.py", 72),
Lianmin Zheng's avatar
Lianmin Zheng committed
147
        TestFile("test_patch_torch.py", 19),
Shangming Cai's avatar
Shangming Cai committed
148
        TestFile("test_release_memory_occupation.py", 257),
Lianmin Zheng's avatar
Lianmin Zheng committed
149
150
    ],
    "per-commit-4-gpu": [
151
152
        TestFile("models/test_qwen3_next_models.py", 291),
        TestFile("test_disaggregation_dp_attention.py", 155),
Lianmin Zheng's avatar
Lianmin Zheng committed
153
154
        TestFile("test_gpt_oss_4gpu.py", 300),
        TestFile("test_local_attn.py", 411),
Lianmin Zheng's avatar
Lianmin Zheng committed
155
        TestFile("test_multi_instance_release_memory_occupation.py", 64),
156
        TestFile("test_pp_single_node.py", 481),
Lianmin Zheng's avatar
Lianmin Zheng committed
157
158
    ],
    "per-commit-8-gpu": [
159
        TestFile("lora/test_lora_llama4.py", 400),
160
161
        TestFile("test_deepseek_v3_basic.py", 275),
        TestFile("test_deepseek_v3_mtp.py", 275),
162
163
        TestFile("test_disaggregation_different_tp.py", 600),
        TestFile("test_disaggregation_pp.py", 140),
Lianmin Zheng's avatar
Lianmin Zheng committed
164
    ],
Yineng Zhang's avatar
Yineng Zhang committed
165
    "per-commit-4-gpu-b200": [
166
        # TestFile("test_gpt_oss_4gpu.py", 600),
167
        # TestFile("test_deepseek_v3_fp4_4gpu.py", 3600),
Lianmin Zheng's avatar
Lianmin Zheng committed
168
169
    ],
    "per-commit-4-gpu-deepep": [
170
        TestFile("ep/test_deepep_small.py", 531),
Lianmin Zheng's avatar
Lianmin Zheng committed
171
172
    ],
    "per-commit-8-gpu-deepep": [
173
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
174
    ],
175
176
177
    "per-commit-8-gpu-deepseek-v32": [
        TestFile("test_deepseek_v32_basic.py", 275),
    ],
178
179
180
    "per-commit-8-gpu-h20": [
        TestFile("quant/test_w4a8_deepseek_v3.py", 371),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
181
    "vllm_dependency_test": [
182
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
183
184
185
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
186
        # TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
187
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
188
189
190
    # If the test cases take too long, considering adding them to nightly tests instead of per-commit tests
    "nightly-1-gpu": [],
    "nightly-8-gpu": [],
Lianmin Zheng's avatar
Lianmin Zheng committed
191
192
193
}

# Add AMD tests
194
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
195
suite_amd = {
196
    "per-commit-amd": [
197
        TestFile("function_call/test_json_schema_constraint.py", 30),
198
199
200
        TestFile("hicache/test_hicache.py", 116),
        TestFile("hicache/test_hicache_mla.py", 127),
        TestFile("hicache/test_hicache_storage.py", 127),
201
        TestFile("lora/test_lora.py", 200),
202
203
        TestFile("lora/test_lora_backend.py", 99),
        TestFile("lora/test_lora_cuda_graph.py", 250),
204
        TestFile("lora/test_lora_eviction.py", 200),
205
        TestFile("lora/test_lora_qwen3.py", 97),
206
        TestFile("lora/test_multi_lora_backend.py", 60),
207
        TestFile("models/test_compressed_tensors_models.py", 42),
208
209
        TestFile("models/test_qwen_models.py", 82),
        TestFile("models/test_reward_models.py", 132),
210
        TestFile("models/test_transformers_models.py", 320),
211
212
        TestFile("openai_server/basic/test_openai_embedding.py", 141),
        TestFile("openai_server/basic/test_openai_server.py", 149),
213
214
215
216
        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),
217
        TestFile("openai_server/features/test_enable_thinking.py", 70),
218
219
220
        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),
221
        TestFile("openai_server/features/test_reasoning_content.py", 89),
222
223
        TestFile("openai_server/function_call/test_openai_function_calling.py", 60),
        TestFile("openai_server/function_call/test_tool_choice.py", 226),
224
        TestFile("openai_server/validation/test_large_max_new_tokens.py", 41),
225
226
        TestFile("openai_server/validation/test_matched_stop.py", 60),
        TestFile("openai_server/validation/test_openai_server_ignore_eos.py", 85),
227
        TestFile("openai_server/validation/test_request_length_validation.py", 31),
228
        TestFile("quant/test_awq_dequant.py", 2),
229
        TestFile("quant/test_block_int8.py", 22),
230
        TestFile("rl/test_update_weights_from_disk.py", 114),
231
232
        TestFile("test_abort.py", 51),
        TestFile("test_chunked_prefill.py", 313),
233
        TestFile("test_create_kvindices.py", 2),
234
        TestFile("test_ebnf_constrained.py", 108),
235
236
        TestFile("test_eval_fp8_accuracy.py", 303),
        TestFile("test_function_call_parser.py", 10),
237
        TestFile("test_fused_moe.py", 30),
238
        TestFile("test_input_embeddings.py", 38),
239
240
241
        TestFile("test_io_struct.py", 8),
        TestFile("test_jinja_template_utils.py", 1),
        TestFile("test_metrics.py", 32),
242
        TestFile("test_metrics_utils.py", 1),
243
244
        TestFile("test_mla.py", 242),
        TestFile("test_mla_deepseek_v3.py", 221),
245
246
        TestFile("test_no_chunked_prefill.py", 108),
        TestFile("test_page_size.py", 60),
247
        TestFile("test_penalty.py", 41),
248
249
        TestFile("test_pytorch_sampling_backend.py", 66),
        TestFile("test_radix_attention.py", 105),
250
        TestFile("test_reasoning_parser.py", 5),
251
        TestFile("test_regex_constrained.py", 64),
252
        TestFile("test_retract_decode.py", 54),
253
        TestFile("test_rope_rocm.py", 3),
254
255
        TestFile("test_server_args.py", 1),
        TestFile("test_skip_tokenizer_init.py", 117),
256
        TestFile("test_srt_endpoint.py", 130),
257
        TestFile("test_srt_engine.py", 261),
258
        TestFile("test_torch_compile.py", 169),
259
        TestFile("test_torch_compile_moe.py", 172),
260
261
        TestFile("test_torch_native_attention_backend.py", 123),
        TestFile("test_triton_attention_backend.py", 150),
262
        TestFile("test_wave_attention_kernels.py", 2),
263
264
265
266
267
268
        # Disabled temporarily
        # TestFile("models/test_embedding_models.py", 73), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
        # TestFile("openai_server/features/test_openai_server_hidden_states.py", 240),
        # TestFile("rl/test_update_weights_from_tensor.py", 48),
        # TestFile("test_no_overlap_scheduler.py", 234), # Disabled temporarily and track in #7703
        # TestFile("test_vision_chunked_prefill.py", 175), # Disabled temporarily and track in #7701
Sai Enduri's avatar
Sai Enduri committed
269
        # TestFile("test_wave_attention_backend.py", 150), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
270
    ],
271
272
    "per-commit-amd-mi35x": [
        TestFile("test_gpt_oss_1gpu.py", 600),
273
        TestFile("test_mla.py", 242),
274
    ],
275
    "per-commit-2-gpu-amd": [
276
277
        TestFile("lora/test_lora_tp.py", 116),
        TestFile("rl/test_update_weights_from_distributed.py", 103),
278
        TestFile("test_data_parallelism.py", 73),
279
        TestFile("test_load_weights_from_remote_instance.py", 72),
Sai Enduri's avatar
Sai Enduri committed
280
        # TestFile("test_patch_torch.py", 19), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
281
    ],
282
283
284
    "per-commit-4-gpu-amd": [
        TestFile("test_pp_single_node.py", 150),
    ],
285
    "per-commit-8-gpu-amd": [
286
287
        TestFile("test_deepseek_v3_basic.py", 275),
        TestFile("test_deepseek_v3_mtp.py", 275),
288
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
289
290
    "nightly-amd": [
        TestFile("test_nightly_gsm8k_eval_amd.py"),
fzyzcjy's avatar
fzyzcjy committed
291
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
292
293
294
}

# Add Intel Xeon tests
295
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
296
suite_xeon = {
297
298
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
299
        TestFile("cpu/test_binding.py"),
300
301
302
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
303
        TestFile("cpu/test_mla.py"),
304
305
306
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
307
        TestFile("cpu/test_rope.py"),
308
        TestFile("cpu/test_shared_expert.py"),
309
        TestFile("cpu/test_topk.py"),
310
        TestFile("test_cpu_graph.py"),
311
        TestFile("test_intel_amx_attention_backend.py"),
312
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
313
314
315
}

# Add Ascend NPU tests
316
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
317
318
suite_ascend = {
    "per-commit-1-ascend-npu": [
319
        TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400),
320
        TestFile("ascend/test_ascend_tp1_bf16.py", 400),
321
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
322
    "per-commit-2-ascend-npu": [
323
        TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),
324
        TestFile("ascend/test_ascend_mla_fia_w8a8int8.py", 400),
325
326
        TestFile("ascend/test_ascend_tp2_bf16.py", 400),
        TestFile("ascend/test_ascend_tp2_fia_bf16.py", 400),
327
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
328
    "per-commit-4-ascend-npu": [
329
        TestFile("ascend/test_ascend_mla_w8a8int8.py", 400),
330
        TestFile("ascend/test_ascend_tp4_bf16.py", 400),
331
    ],
332
333
334
    "per-commit-16-ascend-a3": [
        TestFile("ascend/test_ascend_deepep.py", 400),
    ],
335
336
}

Lianmin Zheng's avatar
Lianmin Zheng committed
337
338
339
340
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)

Lianmin Zheng's avatar
Lianmin Zheng committed
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384

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]

385
386
387
388
389
390

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--timeout-per-file",
        type=int,
391
        default=1200,
392
393
394
395
396
397
398
399
400
        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",
    )
401
402
403
404
405
406
407
408
409
410
411
412
    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
413
414
415
416
417
418
419
420
421
422
    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.",
    )
423
    args = arg_parser.parse_args()
Lianmin Zheng's avatar
Lianmin Zheng committed
424
    print(f"{args=}")
425
426
427
428
429
430

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

Lianmin Zheng's avatar
Lianmin Zheng committed
431
432
433
434
    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]
435

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

438
439
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)