Skip to content

Commit 0e2da4c

Browse files
wangxiyuanwangxiyuan
committed
[UT] refactor test_expert_load_balancer
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
1 parent ebb2a70 commit 0e2da4c

File tree

4 files changed

+218
-153
lines changed

4 files changed

+218
-153
lines changed

‎tests/e2e/singlecard/core/test_ascend_scheduler.py‎

Lines changed: 50 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -98,11 +98,7 @@ def create_scheduler(
9898
)
9999
kv_cache_config = KVCacheConfig(
100100
num_blocks=num_blocks, # A large number of blocks to hold all requests
101-
**({
102-
"tensors": {}
103-
} if vllm_version_is("0.9.0") else {
104-
"kv_cache_tensors": []
105-
}),
101+
kv_cache_tensors=[],
106102
kv_cache_groups=[
107103
KVCacheGroupSpec(['layer'],
108104
FullAttentionSpec(block_size, 1, 1, torch.float32,
@@ -145,8 +141,8 @@ def create_requests(num_requests: int,
145141
multi_modal_hashes=None,
146142
eos_token_id=EOS_TOKEN_ID,
147143
**({
148-
"arrival_time": 0.0
149-
} if vllm_version_is("0.9.0") else {}),
144+
"pooling_params": None
145+
} if not vllm_version_is("0.9.1") else {}),
150146
)
151147
requests.append(request)
152148
return requests
@@ -262,7 +258,9 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
262258
spec_token_ids=None,
263259
logprobs=None,
264260
prompt_logprobs_dict={},
265-
)
261+
**({
262+
"pooler_output": []
263+
} if not vllm_version_is("0.9.1") else {}))
266264
scheduler.update_from_output(output, model_runner_output)
267265

268266
# Schedule the next step. All three requests are running.
@@ -286,7 +284,10 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
286284
spec_token_ids=None,
287285
logprobs=None,
288286
prompt_logprobs_dict={},
289-
)
287+
**({
288+
"pooler_output": []
289+
} if not vllm_version_is("0.9.1") else {}))
290+
290291
scheduler.update_from_output(output1, model_runner_output)
291292
output2 = scheduler.schedule()
292293
assert len(scheduler.running) == 3
@@ -337,7 +338,10 @@ def test_stop_via_update_from_output():
337338
11]], # First request hits EOS, second continues
338339
spec_token_ids=None,
339340
logprobs=None,
340-
prompt_logprobs_dict={})
341+
prompt_logprobs_dict={},
342+
**({
343+
"pooler_output": []
344+
} if not vllm_version_is("0.9.1") else {}))
341345

342346
scheduler.update_from_output(scheduler_output, model_output)
343347

@@ -385,7 +389,10 @@ def test_stop_via_update_from_output():
385389
[13, 14]], # First request hits stop token
386390
spec_token_ids=None,
387391
logprobs=None,
388-
prompt_logprobs_dict={})
392+
prompt_logprobs_dict={},
393+
**({
394+
"pooler_output": []
395+
} if not vllm_version_is("0.9.1") else {}))
389396

390397
scheduler.update_from_output(scheduler_output, model_output)
391398

@@ -432,7 +439,10 @@ def test_stop_via_update_from_output():
432439
[13]], # First request exceeds max_tokens
433440
spec_token_ids=None,
434441
logprobs=None,
435-
prompt_logprobs_dict={})
442+
prompt_logprobs_dict={},
443+
**({
444+
"pooler_output": []
445+
} if not vllm_version_is("0.9.1") else {}))
436446

437447
scheduler.update_from_output(scheduler_output, model_output)
438448

@@ -474,7 +484,10 @@ def test_stop_via_update_from_output():
474484
sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
475485
spec_token_ids=None,
476486
logprobs=None,
477-
prompt_logprobs_dict={})
487+
prompt_logprobs_dict={},
488+
**({
489+
"pooler_output": []
490+
} if not vllm_version_is("0.9.1") else {}))
478491

479492
scheduler.update_from_output(scheduler_output, model_output)
480493

@@ -524,7 +537,10 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],
524537
spec_token_ids=None,
525538
logprobs=None,
526539
prompt_logprobs_dict={},
527-
)
540+
**({
541+
"pooler_output": []
542+
} if not vllm_version_is("0.9.1") else {}))
543+
528544
scheduler.update_from_output(scheduler_output0, model_runner_output)
529545

530546
# Schedule the next step.
@@ -541,7 +557,10 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],
541557
spec_token_ids=None,
542558
logprobs=None,
543559
prompt_logprobs_dict={},
544-
)
560+
**({
561+
"pooler_output": []
562+
} if not vllm_version_is("0.9.1") else {}))
563+
545564
scheduler.update_from_output(scheduler_output1, model_runner_output)
546565

547566

@@ -565,8 +584,6 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
565584
1. Speculated tokens get scheduled correctly
566585
2. Spec decoding stats properly count number of draft and accepted tokens
567586
"""
568-
if vllm_version_is("0.9.0"):
569-
return
570587
num_spec_tokens = max(1, max(len(t) for t in spec_tokens))
571588
scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
572589
requests = create_requests(num_requests=len(spec_tokens), num_tokens=1)
@@ -593,7 +610,10 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
593610
spec_token_ids=spec_tokens,
594611
logprobs=None,
595612
prompt_logprobs_dict={},
596-
)
613+
**({
614+
"pooler_output": []
615+
} if not vllm_version_is("0.9.1") else {}))
616+
597617
engine_core_outputs = scheduler.update_from_output(output,
598618
model_runner_output)
599619

@@ -632,7 +652,10 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
632652
spec_token_ids=None,
633653
logprobs=None,
634654
prompt_logprobs_dict={},
635-
)
655+
**({
656+
"pooler_output": []
657+
} if not vllm_version_is("0.9.1") else {}))
658+
636659
engine_core_outputs = scheduler.update_from_output(output,
637660
model_runner_output)
638661

@@ -727,7 +750,9 @@ def make_output(scheduler: AscendScheduler):
727750
spec_token_ids=None,
728751
logprobs=None,
729752
prompt_logprobs_dict={},
730-
)
753+
**({
754+
"pooler_output": []
755+
} if not vllm_version_is("0.9.1") else {}))
731756

732757

733758
def assert_scheduler_empty(scheduler: AscendScheduler):
@@ -744,11 +769,10 @@ def assert_scheduler_empty(scheduler: AscendScheduler):
744769
assert len(scheduler.encoder_cache_manager.cached) == 0
745770

746771
# KVCache Manager.
747-
if not vllm_version_is("0.9.0"):
748-
assert len(scheduler.kv_cache_manager.coordinator.
749-
single_type_managers[0].req_to_blocks) == 0
750-
assert len(scheduler.kv_cache_manager.coordinator.
751-
single_type_managers[0].num_cached_block) == 0
772+
assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
773+
req_to_blocks) == 0
774+
assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
775+
num_cached_block) == 0
752776
assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
753777
num_free_blocks = (
754778
scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
@@ -789,4 +813,4 @@ def test_memory_leak():
789813
scheduler.update_from_output(scheduler_output, model_runner_output)
790814

791815
# Confirm no memory leak.
792-
assert_scheduler_empty(scheduler)
816+
assert_scheduler_empty(scheduler)

‎tests/e2e/singlecard/test_scheduler.py‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from vllm.v1.structured_output import StructuredOutputManager
3232

3333
from vllm_ascend.core.scheduler import AscendScheduler
34+
from vllm_ascend.utils import vllm_version_is
3435

3536
EOS_TOKEN_ID = 50256
3637

@@ -126,11 +127,14 @@ def create_requests(num_requests: int,
126127
request_id=f"{i}",
127128
prompt_token_ids=[i] * num_tokens,
128129
sampling_params=sampling_params,
130+
pooling_params=None,
129131
multi_modal_inputs=mm_inputs,
130132
multi_modal_placeholders=mm_position,
131133
multi_modal_hashes=None,
132134
eos_token_id=EOS_TOKEN_ID,
133-
)
135+
**({
136+
"pooling_params": None
137+
} if not vllm_version_is("0.9.1") else {}),)
134138
requests.append(request)
135139
return requests
136140

@@ -237,7 +241,10 @@ def test_stop_via_update_from_output():
237241
11]], # First request hits EOS, second continues
238242
spec_token_ids=None,
239243
logprobs=None,
240-
prompt_logprobs_dict={})
244+
prompt_logprobs_dict={},
245+
**({
246+
"pooler_output": []
247+
} if not vllm_version_is("0.9.1") else {}))
241248

242249
scheduler.update_from_output(scheduler_output, model_output)
243250

@@ -283,7 +290,10 @@ def test_stop_via_update_from_output():
283290
[13, 14]], # First request hits stop token
284291
spec_token_ids=None,
285292
logprobs=None,
286-
prompt_logprobs_dict={})
293+
prompt_logprobs_dict={},
294+
**({
295+
"pooler_output": []
296+
} if not vllm_version_is("0.9.1") else {}))
287297

288298
scheduler.update_from_output(scheduler_output, model_output)
289299

@@ -328,7 +338,10 @@ def test_stop_via_update_from_output():
328338
[13]], # First request exceeds max_tokens
329339
spec_token_ids=None,
330340
logprobs=None,
331-
prompt_logprobs_dict={})
341+
prompt_logprobs_dict={},
342+
**({
343+
"pooler_output": []
344+
} if not vllm_version_is("0.9.1") else {}))
332345

333346
scheduler.update_from_output(scheduler_output, model_output)
334347

@@ -369,7 +382,10 @@ def test_stop_via_update_from_output():
369382
sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],
370383
spec_token_ids=None,
371384
logprobs=None,
372-
prompt_logprobs_dict={})
385+
prompt_logprobs_dict={},
386+
**({
387+
"pooler_output": []
388+
} if not vllm_version_is("0.9.1") else {}))
373389

374390
scheduler.update_from_output(scheduler_output, model_output)
375391

0 commit comments

Comments
 (0)