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

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

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

336
337
338
339
340
341
342
# Add Intel XPU tests
suite_xpu = {
    "per-commit-xpu": [
        TestFile("xpu/test_intel_xpu_backend.py"),
    ],
}

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

Lianmin Zheng's avatar
Lianmin Zheng committed
365
366
367
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)
368
suites.update(suite_xpu)
Lianmin Zheng's avatar
Lianmin Zheng committed
369

Lianmin Zheng's avatar
Lianmin Zheng committed
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
410
411
412
413

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]

414
415
416
417
418
419

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

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

Lianmin Zheng's avatar
Lianmin Zheng committed
460
461
462
463
    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]
464

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

467
468
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)