run_suite.py 20.6 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),
Ke Bao's avatar
Ke Bao committed
154
        TestFile("models/test_kimi_linear_models.py", 90),
155
        TestFile("rl/test_update_weights_from_distributed.py", 103),
Lianmin Zheng's avatar
Lianmin Zheng committed
156
        TestFile("test_data_parallelism.py", 73),
157
        TestFile("test_disaggregation_basic.py", 400),
158
        TestFile("test_dp_attention.py", 350),
159
        TestFile("test_load_weights_from_remote_instance.py", 72),
Lianmin Zheng's avatar
Lianmin Zheng committed
160
        TestFile("test_patch_torch.py", 19),
161
        TestFile("test_release_memory_occupation.py", 200),
162
        TestFile("test_eagle_dp_attention.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
163
164
    ],
    "per-commit-4-gpu": [
165
166
        TestFile("models/test_qwen3_next_models.py", 291),
        TestFile("test_disaggregation_dp_attention.py", 155),
Lianmin Zheng's avatar
Lianmin Zheng committed
167
168
        TestFile("test_gpt_oss_4gpu.py", 300),
        TestFile("test_local_attn.py", 411),
Lianmin Zheng's avatar
Lianmin Zheng committed
169
        TestFile("test_multi_instance_release_memory_occupation.py", 64),
170
        TestFile("test_pp_single_node.py", 481),
Lianmin Zheng's avatar
Lianmin Zheng committed
171
    ],
172
    "per-commit-8-gpu-h200": [
173
        TestFile("lora/test_lora_llama4.py", 400),
174
175
        TestFile("test_deepseek_v3_basic.py", 275),
        TestFile("test_deepseek_v3_mtp.py", 275),
176
        TestFile("test_disaggregation_hybrid_attention.py", 200),
Lianmin Zheng's avatar
Lianmin Zheng committed
177
    ],
178
    "per-commit-8-gpu-h20": [
179
        TestFile("quant/test_w4a8_deepseek_v3.py", 520),
180
181
182
        TestFile("test_disaggregation_different_tp.py", 600),
        TestFile("test_disaggregation_pp.py", 140),
    ],
Yineng Zhang's avatar
Yineng Zhang committed
183
    "per-commit-4-gpu-b200": [
184
        TestFile("test_deepseek_v3_fp4_4gpu.py", 3600),
Kangyan-Zhou's avatar
Kangyan-Zhou committed
185
186
        TestFile("test_flash_attention_4.py", 300),
        TestFile("test_gpt_oss_4gpu.py", 600),
187
        TestFile("test_llama31_fp4.py", 300),
Lianmin Zheng's avatar
Lianmin Zheng committed
188
    ],
189
190
191
    "per-commit-4-gpu-gb200": [
        TestFile("test_deepseek_v3_fp4_4gpu.py", 3600),
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
192
    "per-commit-4-gpu-deepep": [
193
        TestFile("ep/test_deepep_small.py", 531),
194
        TestFile("ep/test_mooncake_ep_small.py", 450),
Lianmin Zheng's avatar
Lianmin Zheng committed
195
    ],
196
    "per-commit-8-gpu-h200-deepep": [
197
        TestFile("ep/test_deepep_large.py", 338),
Lianmin Zheng's avatar
Lianmin Zheng committed
198
    ],
199
    "per-commit-8-gpu-h200-deepseek-v32": [
200
        TestFile("test_deepseek_v32_basic.py", 275),
201
        TestFile("test_deepseek_v32_mtp.py", 275),
202
        TestFile("test_deepseek_v32_nsabackend.py", 600),
203
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
204
    "vllm_dependency_test": [
205
        TestFile("quant/test_awq.py", 163),
Lianmin Zheng's avatar
Lianmin Zheng committed
206
207
208
        TestFile("test_bnb.py", 5),
        TestFile("test_gptqmodel_dynamic.py", 102),
        TestFile("test_vllm_dependency.py", 185),
209
        TestFile("test_gguf.py", 96),
Lianmin Zheng's avatar
Lianmin Zheng committed
210
    ],
Lianmin Zheng's avatar
Lianmin Zheng committed
211
212
213
    # 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
214
215
216
}

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

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

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

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

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

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

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]

419
420
421
422
423
424

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

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

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

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

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