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("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),
Lianmin Zheng's avatar
Lianmin Zheng committed
74
        TestFile("test_abort.py", 51),
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", 300),
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_skip_tokenizer_init.py", 117),
Lianmin Zheng's avatar
Lianmin Zheng committed
123
        TestFile("test_srt_endpoint.py", 130),
124
125
        TestFile("test_srt_engine.py", 261),
        TestFile("test_standalone_speculative_decoding.py", 250),
126
        TestFile("test_start_profile.py", 60),
127
128
        TestFile("test_profile_merger.py", 60),
        TestFile("test_profile_merger_http_api.py", 15),
129
        TestFile("test_swa_unittest.py", 1),
Lianmin Zheng's avatar
Lianmin Zheng committed
130
        TestFile("test_torch_compile.py", 76),
Lianmin Zheng's avatar
Lianmin Zheng committed
131
        TestFile("test_torch_compile_moe.py", 172),
132
        TestFile("test_torch_native_attention_backend.py", 123),
Lianmin Zheng's avatar
Lianmin Zheng committed
133
        TestFile("test_torchao.py", 70),
134
        TestFile("test_triton_attention_backend.py", 150),
135
        TestFile("test_triton_attention_kernels.py", 4),
Lianmin Zheng's avatar
Lianmin Zheng committed
136
        TestFile("test_triton_moe_channel_fp8_kernel.py", 25),
137
        TestFile("test_triton_sliding_window.py", 250),
138
        TestFile("test_utils_update_weights.py", 48),
Lianmin Zheng's avatar
Lianmin Zheng committed
139
        TestFile("test_vision_chunked_prefill.py", 175),
140
        TestFile("test_vision_openai_server_a.py", 918),
141
        TestFile("test_vlm_input_format.py", 300),
142
143
        TestFile("test_modelopt_loader.py", 30),
        TestFile("test_modelopt_export.py", 30),
144
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
145
    "per-commit-2-gpu": [
146
        TestFile("ep/test_moe_ep.py", 140),
147
148
149
        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),
150
        TestFile("layers/attention/mamba/test_mamba2_mixer.py", 50),
151
        TestFile("lora/test_lora_tp.py", 116),
152
        TestFile("models/test_glm4_moe_models.py", 100),
153
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
154
        TestFile("test_data_parallelism.py", 73),
155
        TestFile("test_disaggregation_basic.py", 400),
Lianmin Zheng's avatar
Lianmin Zheng committed
156
        TestFile("test_dp_attention.py", 594),
157
        TestFile("test_load_weights_from_remote_instance.py", 72),
Lianmin Zheng's avatar
Lianmin Zheng committed
158
        TestFile("test_patch_torch.py", 19),
Shangming Cai's avatar
Shangming Cai committed
159
        TestFile("test_release_memory_occupation.py", 257),
Lianmin Zheng's avatar
Lianmin Zheng committed
160
161
    ],
    "per-commit-4-gpu": [
162
163
        TestFile("models/test_qwen3_next_models.py", 291),
        TestFile("test_disaggregation_dp_attention.py", 155),
Lianmin Zheng's avatar
Lianmin Zheng committed
164
165
        TestFile("test_gpt_oss_4gpu.py", 300),
        TestFile("test_local_attn.py", 411),
Lianmin Zheng's avatar
Lianmin Zheng committed
166
        TestFile("test_multi_instance_release_memory_occupation.py", 64),
167
        TestFile("test_pp_single_node.py", 481),
Lianmin Zheng's avatar
Lianmin Zheng committed
168
    ],
169
    "per-commit-8-gpu-h200": [
170
        TestFile("lora/test_lora_llama4.py", 400),
171
172
        TestFile("test_deepseek_v3_basic.py", 275),
        TestFile("test_deepseek_v3_mtp.py", 275),
173
        TestFile("test_disaggregation_hybrid_attention.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
174
    ],
175
    "per-commit-8-gpu-h20": [
176
        TestFile("quant/test_w4a8_deepseek_v3.py", 520),
177
178
179
        TestFile("test_disaggregation_different_tp.py", 600),
        TestFile("test_disaggregation_pp.py", 140),
    ],
Yineng Zhang's avatar
Yineng Zhang committed
180
    "per-commit-4-gpu-b200": [
Kangyan-Zhou's avatar
Kangyan-Zhou committed
181
182
183
        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
184
185
    ],
    "per-commit-4-gpu-deepep": [
186
        TestFile("ep/test_deepep_small.py", 531),
187
        TestFile("ep/test_mooncake_ep_small.py", 450),
Lianmin Zheng's avatar
Lianmin Zheng committed
188
    ],
189
    "per-commit-8-gpu-h200-deepep": [
190
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
191
    ],
192
    "per-commit-8-gpu-h200-deepseek-v32": [
193
        TestFile("test_deepseek_v32_basic.py", 275),
194
        TestFile("test_deepseek_v32_mtp.py", 275),
195
        TestFile("test_deepseek_v32_nsabackend.py", 600),
196
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
197
    "vllm_dependency_test": [
198
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
199
200
201
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
202
        TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
203
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
204
205
206
    # 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
207
208
209
}

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

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

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

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

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

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

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]

411
412
413
414
415
416

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

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

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

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

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