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 11, 23:45 UTC
Refactor example prompts fixture (#29854)
Signed-off-by: nwaughac@gmail.com
Chukwuma Nwaugha committed
Dec 5, 2025 at 06:44 UTC
6e865b6a83565c5d661091ec0886403edb171794
1 file changed
+25
-22
tests/conftest.py
+25
-22
@@ -27,7 +27,7 @@ import threading
27
from collections.abc import Generator
28
from contextlib import nullcontext
29
from enum import Enum
30
-from typing import Any, Callable, TypedDict, TypeVar, cast
30
+from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING
31
32
import numpy as np
33
import pytest
@@ -67,6 +67,11 @@ from vllm.transformers_utils.utils import maybe_model_redirect
67
from vllm.utils.collection_utils import is_list_of
68
from vllm.utils.torch_utils import set_default_torch_num_threads
69
70
+if TYPE_CHECKING:
71
+ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
72
+ from transformers.generation.utils import GenerateOutput
73
+
74
+
75
logger = init_logger(__name__)
76
77
_TEST_DIR = os.path.dirname(__file__)
@@ -202,10 +207,7 @@ def dynamo_reset():
207
208
@pytest.fixture
209
def example_prompts() -> list[str]:
205
- prompts = []
206
- for filename in _TEST_PROMPTS:
207
- prompts += _read_prompts(filename)
208
- return prompts
210
+ return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
211
212
213
@pytest.fixture
@@ -224,10 +226,7 @@ class DecoderPromptType(Enum):
226
227
@pytest.fixture
228
def example_long_prompts() -> list[str]:
227
- prompts = []
228
- for filename in _LONG_PROMPTS:
229
- prompts += _read_prompts(filename)
230
- return prompts
229
+ return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
230
231
232
@pytest.fixture(scope="session")
@@ -353,10 +352,13 @@ class HfRunner:
352
trust_remote_code=trust_remote_code,
353
)
354
else:
356
- model = auto_cls.from_pretrained(
357
- model_name,
358
- trust_remote_code=trust_remote_code,
359
- **model_kwargs,
355
+ model = cast(
356
+ nn.Module,
357
+ auto_cls.from_pretrained(
358
+ model_name,
359
+ trust_remote_code=trust_remote_code,
360
+ **model_kwargs,
361
+ ),
362
)
363
364
# in case some unquantized custom models are not in same dtype
@@ -374,10 +376,12 @@ class HfRunner:
376
self.model = model
377
378
if not skip_tokenizer_init:
377
- self.tokenizer = AutoTokenizer.from_pretrained(
378
- model_name,
379
- dtype=dtype,
380
- trust_remote_code=trust_remote_code,
379
+ self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
380
+ AutoTokenizer.from_pretrained(
381
+ model_name,
382
+ dtype=dtype,
383
+ trust_remote_code=trust_remote_code,
384
+ )
385
)
386
387
# don't put this import at the top level
@@ -495,7 +499,7 @@ class HfRunner:
499
500
outputs: list[tuple[list[list[int]], list[str]]] = []
501
for inputs in all_inputs:
498
- output_ids = self.model.generate(
502
+ output_ids: torch.Tensor = self.model.generate(
503
**self.wrap_device(inputs),
504
use_cache=True,
505
**kwargs,
@@ -505,8 +509,7 @@ class HfRunner:
509
skip_special_tokens=True,
510
clean_up_tokenization_spaces=False,
511
)
508
- output_ids = output_ids.cpu().tolist()
509
- outputs.append((output_ids, output_str))
512
+ outputs.append((output_ids.cpu().tolist(), output_str))
513
return outputs
514
515
def generate_greedy(
@@ -574,7 +577,7 @@ class HfRunner:
577
578
all_logprobs: list[list[torch.Tensor]] = []
579
for inputs in all_inputs:
577
- output = self.model.generate(
580
+ output: "GenerateOutput" = self.model.generate(
581
**self.wrap_device(inputs),
582
use_cache=True,
583
do_sample=False,
@@ -656,7 +659,7 @@ class HfRunner:
659
all_output_strs: list[str] = []
660
661
for inputs in all_inputs:
659
- output = self.model.generate(
662
+ output: "GenerateOutput" = self.model.generate(
663
**self.wrap_device(inputs),
664
use_cache=True,
665
do_sample=False,