run_suite.py 20.5 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("debug_utils/test_tensor_dump_forward_hook.py", 15),
18
        TestFile("function_call/test_json_schema_constraint.py", 30),
19
        TestFile("hicache/test_hicache.py", 116),
20
        TestFile("hicache/test_hicache_eagle.py", 150),
21
22
        TestFile("hicache/test_hicache_mla.py", 127),
        TestFile("hicache/test_hicache_storage.py", 127),
23
24
        TestFile("layers/attention/mamba/test_causal_conv1d.py", 25),
        TestFile("layers/attention/mamba/test_mamba_ssm.py", 50),
25
26
27
        TestFile("layers/attention/mamba/test_mamba_ssm_ssd.py", 20),
        TestFile("lora/test_lora.py", 150),
        TestFile("lora/test_lora_eviction.py", 240),
28
        TestFile("lora/test_lora_eviction_policy.py", 200),
29
30
        TestFile("lora/test_lora_backend.py", 99),
        TestFile("lora/test_lora_qwen3.py", 97),
31
        TestFile("lora/test_lora_radix_cache.py", 200),
32
        TestFile("lora/test_lora_openai_api.py", 30),
33
        TestFile("lora/test_lora_update.py", 600),
34
35
36
        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", 160),
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
        TestFile("models/test_vlm_models.py", 741),
46
47
        TestFile("openai_server/basic/test_openai_embedding.py", 79),
        TestFile("openai_server/basic/test_openai_server.py", 270),
48
49
50
51
52
        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),
53
54
55
        TestFile("openai_server/features/test_json_constrained.py", 120),
        TestFile("openai_server/features/test_json_mode.py", 120),
        TestFile("openai_server/features/test_openai_server_ebnf.py", 20),
56
57
        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", 150),
59
        TestFile("openai_server/function_call/test_openai_function_calling.py", 60),
60
        TestFile("openai_server/function_call/test_tool_choice.py", 120),
61
62
63
64
        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", 160),
70
        TestFile("quant/test_autoround.py", 60),
71
        TestFile("rl/test_fp32_lm_head.py", 30),
72
73
74
        TestFile("rl/test_update_weights_from_disk.py", 210),
        TestFile("rl/test_update_weights_from_tensor.py", 80),
        TestFile("test_abort.py", 190),
75
        TestFile("test_build_eagle_tree.py", 8),
76
        TestFile("test_chunked_prefill.py", 410),
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", 500),
81
        TestFile("test_eagle_infer_beta.py", 90),
82
        TestFile("test_ebnf_constrained.py", 80),
Lianmin Zheng's avatar
Lianmin Zheng committed
83
        TestFile("test_eval_fp8_accuracy.py", 303),
84
85
86
        TestFile("test_fa3.py", 420),
        TestFile("test_flashmla.py", 230),
        TestFile("rotary_embedding/test_mrope.py", 10),
87
        TestFile("test_function_call_parser.py", 10),
88
89
        TestFile("test_fused_moe.py", 80),
        TestFile("test_gpt_oss_1gpu.py", 750),
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", 180),
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
        TestFile("test_mla_int8_deepseek_v3.py", 300),
104
        TestFile("test_modelopt_loader.py", 30),
105
        TestFile("test_multi_tokenizer.py", 230),
106
        TestFile("test_ngram_speculative_decoding.py", 290),
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", 82),
112
        TestFile("test_priority_scheduling.py", 130),
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", 450),
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
        TestFile("test_srt_engine.py", 261),
126
        TestFile("test_standalone_speculative_decoding.py", 150),
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),
132
        TestFile("test_torch_compile_moe.py", 210),
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", 100),
139
        TestFile("test_utils_update_weights.py", 48),
140
141
        TestFile("test_vision_chunked_prefill.py", 170),
        TestFile("test_vision_openai_server_a.py", 900),
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
        TestFile("hicache/test_hicache_storage_3fs_backend.py", 200),
        TestFile("hicache/test_hicache_storage_file_backend.py", 200),
150
        TestFile("hicache/test_hicache_storage_mooncake_backend.py", 300),
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),
157
        TestFile("test_dp_attention.py", 350),
158
        TestFile("test_load_weights_from_remote_instance.py", 72),
Lianmin Zheng's avatar
Lianmin Zheng committed
159
        TestFile("test_patch_torch.py", 19),
160
        TestFile("test_release_memory_occupation.py", 200),
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
189
190
    "per-commit-4-gpu-gb200": [
        TestFile("test_deepseek_v3_fp4_4gpu.py", 3600),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
191
    "per-commit-4-gpu-deepep": [
192
        TestFile("ep/test_deepep_small.py", 531),
193
        TestFile("ep/test_mooncake_ep_small.py", 450),
Lianmin Zheng's avatar
Lianmin Zheng committed
194
    ],
195
    "per-commit-8-gpu-h200-deepep": [
196
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
197
    ],
198
    "per-commit-8-gpu-h200-deepseek-v32": [
199
        TestFile("test_deepseek_v32_basic.py", 275),
200
        TestFile("test_deepseek_v32_mtp.py", 275),
201
        TestFile("test_deepseek_v32_nsabackend.py", 600),
202
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
203
    "vllm_dependency_test": [
204
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
205
206
207
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
208
        TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
209
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
210
211
212
    # 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
213
214
215
}

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

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

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

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

Lianmin Zheng's avatar
Lianmin Zheng committed
369
370
371
suites.update(suite_amd)
suites.update(suite_xeon)
suites.update(suite_ascend)
372
suites.update(suite_xpu)
Lianmin Zheng's avatar
Lianmin Zheng committed
373

Lianmin Zheng's avatar
Lianmin Zheng committed
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
414
415
416
417

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]

418
419
420
421
422
423

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

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

Lianmin Zheng's avatar
Lianmin Zheng committed
464
465
466
467
    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]
468

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

471
472
    exit_code = run_unittest_files(files, args.timeout_per_file)
    exit(exit_code)