A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 17:15 UTC
Fix priority preemption regression test in scheduler (#37051)
Signed-off-by: HarshRathva <harshrathvaai@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
HarshRathva committed
Apr 1, 2026 at 09:06 UTC
17b72fd1c8052f7a674056b2228a6e7ebfd4f123
1 file changed
+75
-63
tests/v1/core/test_scheduler.py
+75
-63
@@ -2040,91 +2040,103 @@ def test_priority_scheduling_mixed_priority_and_arrival():
2040
assert scheduled_req_ids == ["3", "2", "1", "0"]
2041
2042
2043
-# This test had previously been passing due to its use of duplicate
2044
-# request ids which resulted in incorrect behavior.
2045
-# Now that the duplicate req ids had been fixed it fails and
2046
-# investigation is needed into whether the priority scheduling
2047
-# preemption logic is working as designed or not.
2048
-@pytest.mark.skip("needs investigation")
2043
def test_priority_scheduling_preemption():
2050
- """Test that priority scheduling preempts
2051
- lower priority requests when memory is constrained."""
2052
- # Create scheduler with very limited memory to force preemption
2044
+ """Test that under KV block pressure the scheduler preempts the
2045
+ lowest-priority *running* request, not the highest-priority one.
2046
+
2047
+ A low-priority request starts running first. Then a high-priority
2048
+ request arrives and is admitted to running. When block pressure
2049
+ builds, the scheduler preempts the low-priority running request
2050
+ while keeping the high-priority one.
2051
+
2052
+ Block math
2053
+ ----------
2054
+ block_size = 16, num_blocks = 6 (1 null → 5 usable).
2055
+
2056
+ Phase 1: lo1 (priority 5, 32 tokens) → 2 blocks. 3 free.
2057
+ Decode → lo1 has 33 tokens (needs 3rd block on next schedule).
2058
+ Phase 2: hi1 (priority 0, 32 tokens) arrives.
2059
+ schedule() allocates lo1's 3rd block (3 used) and admits
2060
+ hi1 (2 blocks) → 5 used, 0 free. Both running.
2061
+ Decode → lo1 34 tokens, hi1 33 tokens.
2062
+ Phase 3: schedule() → hi1 needs 3rd block, 0 free → preemption.
2063
+ lo1 (priority 5) is preempted, hi1 (priority 0) survives.
2064
+ """
2065
+ block_size = 16
2066
+ num_blocks = 6 # 1 null block → 5 usable
2067
+ num_tokens = block_size * 2 # 32 tokens = exactly 2 blocks
2068
+
2069
scheduler = create_scheduler_with_priority(
2054
- max_num_seqs=3, # Allow multiple requests
2070
+ max_num_seqs=3,
2071
max_num_batched_tokens=200,
2056
- num_blocks=6, # Very limited blocks to force memory pressure
2057
- block_size=16, # Standard block size
2058
- )
2059
-
2060
- # Create initial low-priority requests that will consume most memory
2061
- low_priority_requests = create_requests_with_priority(
2062
- num_requests=2,
2063
- priorities=[5, 5], # Low priority
2064
- arrival_times=[1.0, 2.0],
2065
- num_tokens=30, # Large enough to consume significant memory,
2066
- req_ids=["lo1", "lo2"],
2072
+ num_blocks=num_blocks,
2073
+ block_size=block_size,
2074
)
2075
2069
- # Add and schedule low priority requests
2070
- for request in low_priority_requests:
2071
- scheduler.add_request(request)
2076
+ # --- Phase 1: low-priority request starts running ---
2077
+ lo1 = create_requests_with_priority(
2078
+ num_requests=1,
2079
+ priorities=[5],
2080
+ arrival_times=[1.0],
2081
+ num_tokens=num_tokens,
2082
+ req_ids=["lo1"],
2083
+ )[0]
2084
+ scheduler.add_request(lo1)
2085
2086
output = scheduler.schedule()
2074
- assert len(output.scheduled_new_reqs) == 2
2087
+ assert len(output.scheduled_new_reqs) == 1
2088
2076
- # Simulate model execution to move requests to running state
2089
+ # Decode: lo1 now has 33 tokens (crosses 32-token boundary).
2090
model_output = ModelRunnerOutput(
2078
- req_ids=[req.request_id for req in low_priority_requests],
2079
- req_id_to_index={
2080
- req.request_id: i for i, req in enumerate(low_priority_requests)
2081
- },
2082
- sampled_token_ids=[[100] for _ in low_priority_requests],
2091
+ req_ids=["lo1"],
2092
+ req_id_to_index={"lo1": 0},
2093
+ sampled_token_ids=[[100]],
2094
logprobs=None,
2095
prompt_logprobs_dict={},
2096
pooler_output=[],
2097
)
2098
scheduler.update_from_output(output, model_output)
2099
2089
- # Verify both requests are running
2090
- assert len(scheduler.running) == 2
2091
-
2092
- # Now add a high-priority request that requires memory allocation
2093
- # This should trigger preemption due to memory constraints
2094
- high_priority_request = create_requests_with_priority(
2100
+ # --- Phase 2: high-priority request arrives AFTER lo1 is running ---
2101
+ hi1 = create_requests_with_priority(
2102
num_requests=1,
2096
- priorities=[0], # High priority
2097
- arrival_times=[3.0],
2098
- num_tokens=30, # Large enough to require significant memory
2103
+ priorities=[0],
2104
+ arrival_times=[2.0],
2105
+ num_tokens=num_tokens,
2106
req_ids=["hi1"],
2107
)[0]
2108
+ scheduler.add_request(hi1)
2109
2102
- scheduler.add_request(high_priority_request)
2110
+ # schedule(): lo1 gets its 3rd block (3 used), hi1 admitted (5 used,
2111
+ # 0 free). Both are now running.
2112
+ output = scheduler.schedule()
2113
+ assert any(r.req_id == "hi1" for r in output.scheduled_new_reqs)
2114
+ assert len(scheduler.running) == 2
2115
2104
- # Schedule again - this should trigger
2105
- # preemption when trying to allocate memory
2116
+ # Decode: lo1 → 34 tokens, hi1 → 33 tokens.
2117
+ model_output = ModelRunnerOutput(
2118
+ req_ids=["lo1", "hi1"],
2119
+ req_id_to_index={"lo1": 0, "hi1": 1},
2120
+ sampled_token_ids=[[101], [100]],
2121
+ logprobs=None,
2122
+ prompt_logprobs_dict={},
2123
+ pooler_output=[],
2124
+ )
2125
+ scheduler.update_from_output(output, model_output)
2126
+
2127
+ # --- Phase 3: preemption with mixed-priority running requests ---
2128
+ # hi1 needs a 3rd block but 0 are free. The scheduler picks the
2129
+ # lowest-priority running request to preempt:
2130
+ # max(running, key=(priority, arrival_time)) → lo1 (5 > 0).
2131
output = scheduler.schedule()
2132
2108
- # Due to the scheduler's design, if preemption happens
2109
- # during running request scheduling,
2110
- # waiting requests won't be scheduled in the same step
2111
- # Let's check if preemption occurred by looking at the waiting queue
2112
-
2113
- # If preemption happened, we should see requests in the
2114
- # waiting queue
2115
- if len(scheduler.waiting) > 1: # high priority + preempted request
2116
- # Preemption occurred - verify the high priority request
2117
- # gets scheduled next
2118
- output2 = scheduler.schedule()
2119
- assert len(output2.scheduled_new_reqs) == 1
2120
- # High priority request
2121
- assert output2.scheduled_new_reqs[0].req_id == "hi1"
2122
- else:
2123
- # No preemption needed - all requests fit
2124
- # This is also valid behavior if memory allows
2125
- assert len(output.scheduled_new_reqs) == 1
2126
- # High priority request
2127
- assert output.scheduled_new_reqs[0].req_id == "hi1"
2133
+ lo1_req = scheduler.requests["lo1"]
2134
+ assert lo1_req.status == RequestStatus.PREEMPTED, (
2135
+ "Expected low-priority 'lo1' to be preempted"
2136
+ )
2137
+ assert any(req.request_id == "hi1" for req in scheduler.running), (
2138
+ "High-priority 'hi1' should still be running"
2139
+ )
2140
2141
2142
def test_priority_scheduling_no_preemption_when_space_available():