[ModelRunner V2] Update scheduler tests to cover MRV2 paths (#46771)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 26, 2026 at 09:36 UTC
658b54efe419d0e53ec33a8bb7095d8c8b52c741
2 files changed
+53
-18
tests/v1/core/test_scheduler.py
+46
-17
@@ -6,6 +6,7 @@ from unittest.mock import Mock
6
import pytest
7
import torch
8
9
+import vllm.envs as envs
10
from vllm.config import (
11
CacheConfig,
12
ECTransferConfig,
@@ -150,7 +151,7 @@ def test_cached_request_data_resumed_all_token_ids_mrv1_only():
151
"""
152
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
153
153
- scheduler = create_scheduler()
154
+ scheduler = create_scheduler(use_v2_model_runner=False)
155
(req,) = create_requests(num_requests=1, num_tokens=8)
156
req.append_output_token_ids([101, 102, 103])
157
@@ -1899,11 +1900,14 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
1900
assert len(scheduler.waiting) == 0
1901
1902
1903
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
1904
@pytest.mark.parametrize("is_async", [False, True])
1905
@pytest.mark.parametrize(
1906
"use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
1907
)
1906
-def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
1908
+def test_kv_connector_handles_preemption(
1909
+ is_async, use_ec_connector, ec_role, use_v2_model_runner
1910
+):
1911
"""
1912
Test whether scheduler with KVConnector is able to handle
1913
unable to allocate (run out of blocks in allocate_slots().
@@ -1924,6 +1928,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
1928
# encoder connector should not affect test results
1929
use_ec_connector=use_ec_connector,
1930
ec_role=ec_role,
1931
+ use_v2_model_runner=use_v2_model_runner,
1932
)
1933
1934
# Create two requests.
@@ -2034,8 +2039,14 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
2039
)
2040
assert len(scheduler.running) == 1
2041
assert len(scheduler.waiting) == 0
2037
- assert output.scheduled_cached_reqs.num_reqs == 1
2038
- assert output.scheduled_new_reqs == []
2042
+ if use_v2_model_runner:
2043
+ # V2 emits a resumed (previously preempted) request as a
2044
+ # NewRequestData rather than a cached request.
2045
+ assert output.scheduled_cached_reqs.num_reqs == 0
2046
+ assert len(output.scheduled_new_reqs) == 1
2047
+ else:
2048
+ assert output.scheduled_cached_reqs.num_reqs == 1
2049
+ assert output.scheduled_new_reqs == []
2050
_ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
2051
assert len(scheduler.running) == 1
2052
assert len(scheduler.waiting) == 0
@@ -2155,6 +2166,7 @@ def create_scheduler_with_priority(
2166
num_speculative_tokens: int | None = None,
2167
use_ec_connector: bool = False,
2168
ec_role: str | None = None,
2169
+ use_v2_model_runner: bool | None = None,
2170
) -> Scheduler:
2171
"""Create scheduler with priority policy enabled.
2172
@@ -2246,7 +2258,7 @@ def create_scheduler_with_priority(
2258
],
2259
)
2260
cache_config.num_gpu_blocks = num_blocks
2249
- return Scheduler(
2261
+ scheduler = Scheduler(
2262
vllm_config=vllm_config,
2263
kv_cache_config=kv_cache_config,
2264
log_stats=True,
@@ -2254,6 +2266,10 @@ def create_scheduler_with_priority(
2266
block_size=block_size,
2267
hash_block_size=block_size,
2268
)
2269
+ if use_v2_model_runner is None:
2270
+ use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER)
2271
+ scheduler.use_v2_model_runner = use_v2_model_runner
2272
+ return scheduler
2273
2274
2275
_none_hash_initialized = False
@@ -2955,11 +2971,12 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
2971
assert engine_core_output.finish_reason == FinishReason.ERROR
2972
2973
2974
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
2975
@pytest.mark.parametrize(
2976
"use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
2977
)
2978
def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
2962
- use_ec_connector, ec_role
2979
+ use_ec_connector, ec_role, use_v2_model_runner
2980
):
2981
"""Test that priority scheduling preempts lower priority requests
2982
when out of KV cache space."""
@@ -2973,6 +2990,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
2990
# encoder connector should not affect test results
2991
use_ec_connector=use_ec_connector,
2992
ec_role=ec_role,
2993
+ use_v2_model_runner=use_v2_model_runner,
2994
)
2995
2996
# Create a request and schedule it
@@ -3065,20 +3083,31 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
3083
output = scheduler.schedule()
3084
scheduled_cached_reqs = output.scheduled_cached_reqs
3085
3068
- assert len(output.scheduled_new_reqs) == 0
3069
- assert scheduled_cached_reqs.num_reqs == 1
3086
assert len(scheduler.waiting) == 0
3087
assert len(scheduler.running) == 1
3088
3073
- # Preempted request resumed in scheduled_cached_reqs
3074
- assert len(scheduled_cached_reqs.resumed_req_ids) == 1
3075
- assert len(scheduled_cached_reqs.all_token_ids) == 1
3076
- assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
3077
- assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
3078
- assert request_low.request_id in scheduled_cached_reqs.all_token_ids
3079
- # Resumed tokens include 30 prompt tokens and 2 decoded tokens
3080
- assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32
3081
- assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100
3089
+ if use_v2_model_runner:
3090
+ # V2 emits the resumed request as a NewRequestData, carrying its full
3091
+ # token ids in prefill_token_ids (instead of cached all_token_ids).
3092
+ assert scheduled_cached_reqs.num_reqs == 0
3093
+ assert len(output.scheduled_new_reqs) == 1
3094
+ new_req = output.scheduled_new_reqs[0]
3095
+ assert new_req.req_id == request_low.request_id
3096
+ # Resumed tokens include 30 prompt tokens and 2 decoded tokens.
3097
+ assert len(new_req.prefill_token_ids) == 32
3098
+ assert new_req.prefill_token_ids[31] == 100
3099
+ else:
3100
+ assert len(output.scheduled_new_reqs) == 0
3101
+ assert scheduled_cached_reqs.num_reqs == 1
3102
+ # Preempted request resumed in scheduled_cached_reqs
3103
+ assert len(scheduled_cached_reqs.resumed_req_ids) == 1
3104
+ assert len(scheduled_cached_reqs.all_token_ids) == 1
3105
+ assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
3106
+ assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
3107
+ assert request_low.request_id in scheduled_cached_reqs.all_token_ids
3108
+ # Resumed tokens include 30 prompt tokens and 2 decoded tokens
3109
+ assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32
3110
+ assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100
3111
3112
3113
@pytest.mark.parametrize(
tests/v1/core/utils.py
+7
-1
@@ -3,6 +3,7 @@
3
4
import torch
5
6
+import vllm.envs as envs
7
from tests.v1.kv_connector.unit.utils import MockKVConfig
8
from vllm.config import (
9
CacheConfig,
@@ -58,6 +59,7 @@ def create_scheduler(
59
pipeline_parallel_size: int = 1,
60
use_ec_connector: bool = False,
61
ec_role: str | None = None,
62
+ use_v2_model_runner: bool | None = None,
63
) -> Scheduler | AsyncScheduler:
64
"""Create scheduler under test.
65
@@ -165,13 +167,17 @@ def create_scheduler(
167
cache_config.num_gpu_blocks = num_blocks
168
register_all_kvcache_specs(vllm_config)
169
scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
168
- return scheduler_cls(
170
+ scheduler = scheduler_cls(
171
vllm_config=vllm_config,
172
kv_cache_config=kv_cache_config,
173
block_size=block_size,
174
log_stats=True,
175
structured_output_manager=StructuredOutputManager(vllm_config),
176
)
177
+ if use_v2_model_runner is None:
178
+ use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER)
179
+ scheduler.use_v2_model_runner = use_v2_model_runner
180
+ return scheduler
181
182
183
_none_hash_initialized = False