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 15, 17:15 UTC
[TPU][CI] Fix TPUModelRunner Test (#15667)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Mar 28, 2025 at 03:01 UTC
2d9045fce8a6b440f937a5f313bf8bc5baf3103a
2 files changed
+2
-18
.buildkite/run-tpu-v1-test.sh
+1
-1
@@ -30,7 +30,7 @@ docker run --privileged --net host --shm-size=16G -it \
30
&& echo TEST_4 \
31
&& python3 /workspace/vllm/examples/offline_inference/tpu.py \
32
&& echo TEST_5 \
33
- && pytest -s -v /workspace/vllm/tests/tpu/worker/test_tpu_model_runner.py \
33
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py \
34
&& echo TEST_6 \
35
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" \
36
tests/v1/tpu/worker/test_tpu_model_runner.py
+1
-17
@@ -7,7 +7,6 @@ from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
7
from vllm.sampling_params import SamplingParams
8
from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
9
SchedulerOutput)
10
-from vllm.v1.sample.metadata import SamplingMetadata
10
from vllm.v1.worker.tpu_model_runner import (TPUModelRunner,
11
_get_padded_token_len,
12
_get_paddings)
@@ -113,12 +112,6 @@ def _is_req_added(model_runner, req_id: str) -> bool:
112
return req_id in model_runner.requests
113
114
116
-def _is_sampling_metadata_changed(model_runner,
117
- sampling_metadata_before: SamplingMetadata):
118
- return model_runner.input_batch.sampling_metadata is not (
119
- sampling_metadata_before)
120
-
121
-
115
def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
116
req_index = model_runner.input_batch.req_id_to_index[req_id]
117
block_table = model_runner.input_batch.block_table
@@ -136,10 +129,8 @@ def test_update_states_new_request(model_runner):
129
# new req
130
scheduler_output = _schedule_new_request(req_id)
131
139
- metadata_before = model_runner.input_batch.sampling_metadata
132
model_runner._update_states(scheduler_output)
133
142
- assert _is_sampling_metadata_changed(model_runner, metadata_before)
134
assert _is_req_added(model_runner, req_id)
135
assert _is_req_scheduled(model_runner, req_id)
136
assert _is_req_state_block_table_match(model_runner, req_id)
@@ -170,9 +161,7 @@ def test_update_states_request_finished(model_runner):
161
grammar_bitmask=None,
162
)
163
173
- metadata_before = model_runner.input_batch.sampling_metadata
164
model_runner._update_states(scheduler_output)
175
- assert _is_sampling_metadata_changed(model_runner, metadata_before)
165
assert not _is_req_added(model_runner, req_id)
166
assert not _is_req_scheduled(model_runner, req_id)
167
@@ -229,9 +218,7 @@ def test_update_states_request_resumed(model_runner):
218
grammar_bitmask=None,
219
)
220
232
- metadata_before = model_runner.input_batch.sampling_metadata
221
model_runner._update_states(scheduler_output)
234
- assert _is_sampling_metadata_changed(model_runner, metadata_before)
222
assert _is_req_added(model_runner, req_id)
223
assert _is_req_scheduled(model_runner, req_id)
224
assert _is_req_state_block_table_match(model_runner, req_id)
@@ -262,9 +249,7 @@ def test_update_states_no_changes(model_runner):
249
grammar_bitmask=None,
250
)
251
265
- metadata_before = model_runner.input_batch.sampling_metadata
252
model_runner._update_states(scheduler_output)
267
- assert not _is_sampling_metadata_changed(model_runner, metadata_before)
253
assert _is_req_added(model_runner, req_id)
254
assert _is_req_scheduled(model_runner, req_id)
255
assert _is_req_state_block_table_match(model_runner, req_id)
@@ -299,8 +284,7 @@ def test_update_states_request_unscheduled(model_runner):
284
grammar_bitmask=None,
285
)
286
302
- metadata_before = model_runner._update_states(scheduler_output)
303
- assert _is_sampling_metadata_changed(model_runner, metadata_before)
287
+ model_runner._update_states(scheduler_output)
288
289
assert _is_req_added(model_runner, req_ids[0])
290
assert _is_req_scheduled(model_runner, req_ids[0])