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

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

# Add Intel Xeon tests
302
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
303
suite_xeon = {
304
305
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
306
        TestFile("cpu/test_binding.py"),
307
308
309
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
310
        TestFile("cpu/test_mla.py"),
311
312
313
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
314
        TestFile("cpu/test_rope.py"),
315
        TestFile("cpu/test_shared_expert.py"),
316
        TestFile("cpu/test_topk.py"),
317
318
319
320
        TestFile("cpu/test_cpu_graph.py"),
        TestFile("cpu/test_intel_amx_attention_backend_a.py"),
        TestFile("cpu/test_intel_amx_attention_backend_b.py"),
        TestFile("cpu/test_intel_amx_attention_backend_c.py"),
321
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
322
323
}

324
325
326
327
328
329
330
# Add Intel XPU tests
suite_xpu = {
    "per-commit-xpu": [
        TestFile("xpu/test_intel_xpu_backend.py"),
    ],
}

Lianmin Zheng's avatar
Lianmin Zheng committed
331
# Add Ascend NPU tests
332
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
333
334
suite_ascend = {
    "per-commit-1-ascend-npu": [
335
        TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400),
336
        TestFile("ascend/test_ascend_tp1_bf16.py", 400),
337
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
338
    "per-commit-2-ascend-npu": [
339
        TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),
340
        TestFile("ascend/test_ascend_mla_fia_w8a8int8.py", 400),
341
342
        TestFile("ascend/test_ascend_tp2_bf16.py", 400),
        TestFile("ascend/test_ascend_tp2_fia_bf16.py", 400),
343
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
344
    "per-commit-4-ascend-npu": [
345
        TestFile("ascend/test_ascend_mla_w8a8int8.py", 400),
346
        TestFile("ascend/test_ascend_tp4_bf16.py", 400),
347
    ],
348
349
350
    "per-commit-16-ascend-a3": [
        TestFile("ascend/test_ascend_deepep.py", 400),
    ],
351
352
}

Lianmin Zheng's avatar
Lianmin Zheng committed
353
354
355
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)
356
suites.update(suite_xpu)
Lianmin Zheng's avatar
Lianmin Zheng committed
357

Lianmin Zheng's avatar
Lianmin Zheng committed
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
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401

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]

402
403
404
405
406
407

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--timeout-per-file",
        type=int,
408
        default=1200,
409
410
411
412
413
414
415
416
417
        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",
    )
418
419
420
421
422
423
424
425
426
427
428
429
    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
430
431
432
433
434
435
436
437
438
439
    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.",
    )
440
    args = arg_parser.parse_args()
Lianmin Zheng's avatar
Lianmin Zheng committed
441
    print(f"{args=}")
442
443
444
445
446
447

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

Lianmin Zheng's avatar
Lianmin Zheng committed
448
449
450
451
    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]
452

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

455
456
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)