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

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

# Add Intel Xeon tests
310
# NOTE: please sort the test cases alphabetically by the test file name
Lianmin Zheng's avatar
Lianmin Zheng committed
311
suite_xeon = {
312
313
    "per-commit-cpu": [
        TestFile("cpu/test_activation.py"),
314
        TestFile("cpu/test_binding.py"),
315
316
317
        TestFile("cpu/test_decode.py"),
        TestFile("cpu/test_extend.py"),
        TestFile("cpu/test_gemm.py"),
318
        TestFile("cpu/test_mla.py"),
319
320
321
        TestFile("cpu/test_moe.py"),
        TestFile("cpu/test_norm.py"),
        TestFile("cpu/test_qkv_proj_with_rope.py"),
322
        TestFile("cpu/test_rope.py"),
323
        TestFile("cpu/test_shared_expert.py"),
324
        TestFile("cpu/test_topk.py"),
325
326
327
328
        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"),
329
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
330
331
}

332
333
334
335
336
337
338
# Add Intel XPU tests
suite_xpu = {
    "per-commit-xpu": [
        TestFile("xpu/test_intel_xpu_backend.py"),
    ],
}

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

Lianmin Zheng's avatar
Lianmin Zheng committed
361
362
363
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)
364
suites.update(suite_xpu)
Lianmin Zheng's avatar
Lianmin Zheng committed
365

Lianmin Zheng's avatar
Lianmin Zheng committed
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
402
403
404
405
406
407
408
409

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]

410
411
412
413
414
415

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

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

Lianmin Zheng's avatar
Lianmin Zheng committed
456
457
458
459
    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]
460

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

463
464
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)