conftest.py 40.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
import os
5
import tempfile
6
from enum import Enum
7
from typing import Any, Callable, Optional, TypedDict, TypeVar, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
8

9
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
10
11
import pytest
import torch
12
import torch.nn as nn
13
import torch.nn.functional as F
14
from huggingface_hub import snapshot_download
15
from PIL import Image
16
17
from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
                          BatchEncoding, BatchFeature)
18
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
19

20
21
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
22
from vllm import LLM, SamplingParams
23
from vllm.assets.audio import AudioAsset
24
from vllm.assets.image import ImageAsset
25
from vllm.assets.video import VideoAsset
26
from vllm.config import TaskOption, _get_and_verify_dtype
27
from vllm.connections import global_http_connection
28
from vllm.distributed import (cleanup_dist_env_and_memory,
29
30
                              init_distributed_environment,
                              initialize_model_parallel)
31
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
32
                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
33
from vllm.logger import init_logger
34
from vllm.outputs import RequestOutput
35
from vllm.sampling_params import BeamSearchParams
36
from vllm.utils import cuda_device_count_stateless
37

38
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
39

40
41
42
_TEST_DIR = os.path.dirname(__file__)
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
_LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
43
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
44

Cyrus Leung's avatar
Cyrus Leung committed
45
_M = TypeVar("_M")
46

47
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
48
49

PromptImageInput = _PromptMultiModalInput[Image.Image]
50
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
51
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
52

53

54
def _read_prompts(filename: str) -> list[str]:
55
    with open(filename) as f:
56
57
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
58
59


60
class ImageAssetPrompts(TypedDict):
61
62
    stop_sign: str
    cherry_blossom: str
63
64


65
class ImageTestAssets(list[ImageAsset]):
66
67

    def __init__(self) -> None:
68
69
70
71
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
72

73
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
74
75
76
77
78
79
        """
        Convenience method to define the prompt for each test image.

        The order of the returned prompts matches the order of the
        assets when iterating through this object.
        """
80
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
81
82


83
84
class VideoAssetPrompts(TypedDict):
    baby_reading: str
85
86


87
class VideoTestAssets(list[VideoAsset]):
88
89
90

    def __init__(self) -> None:
        super().__init__([
91
            VideoAsset("baby_reading"),
92
93
        ])

94
95
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
96
97


98
class AudioAssetPrompts(TypedDict):
99
100
101
102
    mary_had_lamb: str
    winning_call: str


103
class AudioTestAssets(list[AudioAsset]):
104
105
106
107
108
109
110

    def __init__(self) -> None:
        super().__init__([
            AudioAsset("mary_had_lamb"),
            AudioAsset("winning_call"),
        ])

111
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
112
113
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

114

115
IMAGE_ASSETS = ImageTestAssets()
116
"""Singleton instance of {class}`ImageTestAssets`."""
117
VIDEO_ASSETS = VideoTestAssets()
118
"""Singleton instance of {class}`VideoTestAssets`."""
119
AUDIO_ASSETS = AudioTestAssets()
120
"""Singleton instance of {class}`AudioTestAssets`."""
121
122


123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@pytest.fixture(scope="function", autouse=True)
def cleanup_VLLM_USE_V1(monkeypatch):
    """
    The V1 oracle sets "VLLM_USE_V1" during loading. This means
    that each invocation of a test change the env variable.

    If we touch "VLLM_USE_V1" with monkeypatch, then any changes
    made during the test run by vLLM will be cleaned up.

    This fixture is used by every test.
    """

    # If VLLM_USE_V1 is not set, set then delete. This will
    # cause monkeypatch to clean up VLLM_USE_V1 upon exit
    # if VLLM modifies the value of envs.VLLM_USE_V1.
    if "VLLM_USE_V1" not in os.environ:
        monkeypatch.setenv("VLLM_USE_V1", "")
        monkeypatch.delenv("VLLM_USE_V1")


Joe Runde's avatar
Joe Runde committed
143
@pytest.fixture(params=[True, False])
144
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
145
146
147
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
148
    skip_v0 = request.node.get_closest_marker("skip_v0")
Joe Runde's avatar
Joe Runde committed
149
150
151
152
153
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
154
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
155
    else:
156
157
        if skip_v0:
            pytest.skip("Skipping test on vllm V0")
158
159
160
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
161
162


163
164
165
166
167
168
169
@pytest.fixture(autouse=True)
def init_test_http_connection():
    # pytest_asyncio may use a different event loop per test
    # so we need to make sure the async client is created anew
    global_http_connection.reuse_client = False


170
171
172
173
174
175
176
177
178
179
180
181
@pytest.fixture
def dist_init():
    temp_file = tempfile.mkstemp()[1]
    init_distributed_environment(
        world_size=1,
        rank=0,
        distributed_init_method=f"file://{temp_file}",
        local_rank=0,
        backend="nccl",
    )
    initialize_model_parallel(1, 1)
    yield
182
    cleanup_dist_env_and_memory()
183
184


185
@pytest.fixture()
186
def should_do_global_cleanup_after_test(request) -> bool:
187
188
189
190
    """Allow subdirectories to skip global cleanup by overriding this fixture.
    This can provide a ~10x speedup for non-GPU unit tests since they don't need
    to initialize torch.
    """
191

192
    return not request.node.get_closest_marker("skip_global_cleanup")
193
194


195
@pytest.fixture(autouse=True)
196
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
197
    yield
198
    if should_do_global_cleanup_after_test:
199
        cleanup_dist_env_and_memory()
200
201


202
203
204
205
206
207
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
208
@pytest.fixture
209
def example_prompts() -> list[str]:
210
211
    prompts = []
    for filename in _TEST_PROMPTS:
212
        prompts += _read_prompts(filename)
213
214
215
    return prompts


216
217
218
219
220
221
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


222
223
224
225
226
227
228
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


229
@pytest.fixture
230
def example_encoder_decoder_prompts(
231
) -> dict[DecoderPromptType, list[ExplicitEncoderDecoderPrompt]]:
232
233
234
235
236
237
    '''
    Returns an encoder prompt list and a decoder prompt list, wherein each pair
    of same-index entries in both lists corresponds to an (encoder prompt,
    decoder prompt) tuple.

    Returns:
238

239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    * Encoder prompt list
    * Decoder prompt list (reverse of encoder prompt list)
    '''

    encoder_prompts = []
    for filename in _TEST_PROMPTS:
        encoder_prompts += _read_prompts(filename)

    custom_decoder_prompts = encoder_prompts[::-1]
    empty_str_decoder_prompts = [""] * len(encoder_prompts)
    none_decoder_prompts = [None] * len(encoder_prompts)

    # NONE decoder prompt type
    return {
        DecoderPromptType.NONE:
254
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
255
        DecoderPromptType.EMPTY_STR:
256
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
257
        DecoderPromptType.CUSTOM:
258
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
259
260
261
    }


262
@pytest.fixture
263
def example_long_prompts() -> list[str]:
264
265
    prompts = []
    for filename in _LONG_PROMPTS:
266
        prompts += _read_prompts(filename)
267
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
268
269


270
@pytest.fixture(scope="session")
271
def image_assets() -> ImageTestAssets:
272
273
274
    return IMAGE_ASSETS


275
@pytest.fixture(scope="session")
276
def video_assets() -> VideoTestAssets:
277
278
279
    return VIDEO_ASSETS


280
@pytest.fixture(scope="session")
281
def audio_assets() -> AudioTestAssets:
282
283
284
    return AUDIO_ASSETS


285
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
286
_R = TypeVar("_R")
287

Woosuk Kwon's avatar
Woosuk Kwon committed
288
289
290

class HfRunner:

291
    def get_default_device(self):
292
        from vllm.platforms import current_platform
293

294
295
        return ("cpu"
                if current_platform.is_cpu() else current_platform.device_type)
296
297

    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
298
299
300
        if x is None or isinstance(x, (bool, )):
            return x

301
        if device is None:
302
            device = self.device
303

304
305
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
306

307
308
309
310
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
311

Woosuk Kwon's avatar
Woosuk Kwon committed
312
313
314
    def __init__(
        self,
        model_name: str,
315
        dtype: str = "auto",
316
        *,
317
        model_kwargs: Optional[dict[str, Any]] = None,
318
        trust_remote_code: bool = True,
319
        is_sentence_transformer: bool = False,
320
        is_cross_encoder: bool = False,
321
        skip_tokenizer_init: bool = False,
322
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
Woosuk Kwon's avatar
Woosuk Kwon committed
323
    ) -> None:
324
        self.model_name = model_name
325

326
327
        self.config = AutoConfig.from_pretrained(
            model_name,
328
            trust_remote_code=trust_remote_code,
329
330
        )
        self.device = self.get_default_device()
331
332
333
334
335
336
        self.dtype = torch_dtype = _get_and_verify_dtype(
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
337
338
339
340

        model_kwargs = model_kwargs if model_kwargs is not None else {}
        model_kwargs.setdefault("torch_dtype", torch_dtype)

341
        if is_sentence_transformer:
342
343
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
344
345
346
347
348

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
349
                trust_remote_code=trust_remote_code,
350
            )
351
352
353
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
354
355
356
357
358

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
359
                trust_remote_code=trust_remote_code,
360
            )
361
        else:
362
363
            model = auto_cls.from_pretrained(
                model_name,
364
                trust_remote_code=trust_remote_code,
365
366
367
                **model_kwargs,
            )

368
369
370
371
372
373
            # in case some unquantized custom models are not in same dtype
            if (getattr(model, "quantization_method", None) is None
                    and any(p.dtype != self.dtype
                            for p in model.parameters())):
                model = model.to(dtype=self.dtype)

374
375
376
            if (getattr(model, "quantization_method", None) != "bitsandbytes"
                    and len({p.device
                             for p in model.parameters()}) < 2):
377
                model = model.to(device=self.device)
378
379

            self.model = model
380

381
382
383
384
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
385
                trust_remote_code=trust_remote_code,
386
            )
387

388
389
390
391
392
393
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
394
            trust_remote_code=trust_remote_code,
395
        )
396
397
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
398

399
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
400
        self,
401
        prompts: list[str],
402
        images: Optional[PromptImageInput] = None,
403
404
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
405
    ) -> list[Union[BatchFeature, BatchEncoding]]:
406
        if images is not None:
407
            assert len(prompts) == len(images)
408

409
410
411
412
413
414
        if videos is not None:
            assert len(prompts) == len(videos)

        if audios is not None:
            assert len(prompts) == len(audios)

415
        all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
416
        for i, prompt in enumerate(prompts):
417
            processor_kwargs: dict[str, Any] = {
418
419
420
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
421
422
423
424
            if images is not None and (image := images[i]) is not None:
                processor_kwargs["images"] = image
            if videos is not None and (video := videos[i]) is not None:
                processor_kwargs["videos"] = video
425
426
427
428
429
430
431
432
433
            if audios is not None and (audio_inputs := audios[i]) is not None:
                # HACK - not all processors take sampling_rate; we should
                # clean this up in the future.
                if len(audio_inputs) == 2:
                    audio, sr = audio_inputs
                    processor_kwargs["audio"] = audio
                    processor_kwargs["sampling_rate"] = sr
                else:
                    processor_kwargs["audio"] = audio_inputs
434
435

            inputs = self.processor(**processor_kwargs)
436
437
            if isinstance(inputs, BatchFeature):
                inputs = inputs.to(dtype=self.dtype)
438

439
440
441
442
            all_inputs.append(inputs)

        return all_inputs

443
444
445
446
447
448
449
450
451
    def get_prompt_embeddings(self, prompts: list[str]) -> list[torch.Tensor]:
        all_inputs = self.get_inputs(prompts)
        embeddings = []
        for inputs in all_inputs:
            input_ids = self.wrap_device(inputs)["input_ids"]
            embedding = self.model.get_input_embeddings()(input_ids).squeeze(0)
            embeddings.append(embedding)
        return embeddings

452
    def classify(self, prompts: list[str]) -> list[str]:
453
454
455
456
457
458
459
460
461
462
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
            logits = output.logits.softmax(dim=-1)[0].tolist()
            outputs.append(logits)

        return outputs

463
464
    def generate(
        self,
465
        prompts: list[str],
466
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
467
        videos: Optional[PromptVideoInput] = None,
468
469
        audios: Optional[PromptAudioInput] = None,
        **kwargs: Any,
470
    ) -> list[tuple[list[list[int]], list[str]]]:
471
472
473
474
475
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

476
        outputs: list[tuple[list[list[int]], list[str]]] = []
477
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
478
            output_ids = self.model.generate(
479
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
480
481
482
                use_cache=True,
                **kwargs,
            )
483
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
484
485
486
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
487
488
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
489
490
491
492
493
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
494
        prompts: list[str],
Woosuk Kwon's avatar
Woosuk Kwon committed
495
        max_tokens: int,
496
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
497
        videos: Optional[PromptVideoInput] = None,
498
        audios: Optional[PromptAudioInput] = None,
499
        **kwargs: Any,
500
    ) -> list[tuple[list[int], str]]:
501
502
        outputs = self.generate(prompts,
                                do_sample=False,
503
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
504
                                images=images,
505
506
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
507
                                **kwargs)
508
509
510

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
511
512
513

    def generate_beam_search(
        self,
514
        prompts: list[str],
515
516
        beam_width: int,
        max_tokens: int,
517
518
519
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
520
    ) -> list[tuple[list[list[int]], list[str]]]:
521
522
523
524
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
525
526
527
528
529
                                num_return_sequences=beam_width,
                                images=images,
                                videos=videos,
                                audios=audios)

530
531
532
533
534
535
536
537
538
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
                    x for x in output_ids[j]
                    if x != self.tokenizer.pad_token_id
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
539

540
541
    def generate_greedy_logprobs(
        self,
542
        prompts: list[str],
543
        max_tokens: int,
544
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
545
        videos: Optional[PromptVideoInput] = None,
546
        audios: Optional[PromptAudioInput] = None,
547
        **kwargs: Any,
548
    ) -> list[list[torch.Tensor]]:
549
550
551
552
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
553

554
        all_logprobs: list[list[torch.Tensor]] = []
555
        for inputs in all_inputs:
556
            output = self.model.generate(
557
                **self.wrap_device(inputs),
558
559
560
561
562
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
563
                **kwargs,
564
            )
565
566
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
567
568
569
            all_logprobs.append(seq_logprobs)
        return all_logprobs

570
    def _hidden_states_to_seq_logprobs(
571
        self,
572
573
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
574
575
        output_embeddings = self.model.get_output_embeddings()

576
        seq_logprobs: list[torch.Tensor] = []
577
578
579
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
580
581
582
583
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
584
                output_embeddings.weight.t(),
585
            )
586
587
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
588
589
590
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

591
592
593
594
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
595
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
596
        num_logprobs: int,
597
    ) -> tuple[list[dict[int, float]], int]:
598
599
600
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

601
        # convert to dict
602
        seq_logprobs_lst: list[dict[int, float]] = []
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
        for tok_idx, tok_logprobs in enumerate(seq_logprobs):
            # drop prompt logprobs
            if tok_idx == 0:
                tok_logprobs = tok_logprobs[-1, :].reshape(1, -1)
            topk = tok_logprobs.topk(num_logprobs)

            tok_logprobs_dct = {}
            for token_id, logprob in zip(topk.indices[0], topk.values[0]):
                tok_logprobs_dct[token_id.item()] = logprob.item()

            seq_logprobs_lst.append(tok_logprobs_dct)

        return (
            seq_logprobs_lst,
            output_len,
        )

620
621
    def generate_greedy_logprobs_limit(
        self,
622
        prompts: list[str],
623
624
        max_tokens: int,
        num_logprobs: int,
625
626
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
627
        videos: Optional[PromptVideoInput] = None,
628
        **kwargs: Any,
629
    ) -> list[TokensTextLogprobs]:
630
631
632
633
634
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

635
636
637
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
638

639
        for inputs in all_inputs:
640
            output = self.model.generate(
641
                **self.wrap_device(inputs),
642
643
644
645
646
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
647
                **kwargs,
648
649
            )

650
651
652
653
654
655
656
657
658
659
660
661
            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.hidden_states,
                                                num_logprobs)

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_len = len(seq_logprobs_lst)
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))
662

663
664
665
666
667
668
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

    def generate_encoder_decoder_greedy_logprobs_limit(
        self,
669
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
670
671
        max_tokens: int,
        num_logprobs: int,
672
        images: Optional[PromptImageInput] = None,
673
        **kwargs: Any,
674
    ) -> list[TokensTextLogprobs]:
675
676
677
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
678

679
680
681
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
682

683
684
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
685
            processor_kwargs: dict[str, Any] = {
686
687
688
689
690
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
691

692
693
            encoder_inputs = self.processor(**processor_kwargs)
            encoder_inputs = self.wrap_device(encoder_inputs)
694
695
696
697

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
698
699
700
                decoder_inputs = self.tokenizer(decoder_prompt,
                                                return_tensors="pt")
                decoder_input_ids = self.wrap_device(decoder_inputs.input_ids)
701
702
703
704
705
706
707
708

            output = self.model.generate(
                decoder_input_ids=decoder_input_ids,
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
709
                **encoder_inputs,
710
711
712
713
714
715
716
717
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
718
719
720
721
722
723
724
725
726
727
728

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))

        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

729
730
731
    def encode(self, prompts: list[str], *args,
               **kwargs) -> list[list[torch.Tensor]]:
        return self.model.encode(prompts, *args, **kwargs)
732

733
734
735
736
737
738
    def predict(self, prompts: list[list[str]], *args,
                **kwargs) -> torch.Tensor:
        return self.model.predict(prompts,
                                  *args,
                                  convert_to_tensor=True,
                                  **kwargs)
739

740
741
742
743
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
744
        del self.model
745
        cleanup_dist_env_and_memory()
746

Woosuk Kwon's avatar
Woosuk Kwon committed
747

Cyrus Leung's avatar
Cyrus Leung committed
748
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
749
750
751
752
753
def hf_runner():
    return HfRunner


class VllmRunner:
754
755
    """
    The default value of some arguments have been modified from
756
    {class}`~vllm.LLM` as follows:
757

758
759
760
761
762
763
    - `trust_remote_code`: Set to `True` instead of `False` for convenience.
    - `seed`: Set to `0` instead of `None` for test reproducibility.
    - `max_model_len`: Set to `1024` instead of `None` to reduce memory usage.
    - `block_size`: Set to `16` instead of `None` to reduce memory usage.
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
764
    - `enforce_eager`: Set to `False` to test CUDA graph.
765
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
766
767
768
769

    def __init__(
        self,
        model_name: str,
770
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
771
        tokenizer_name: Optional[str] = None,
772
        tokenizer_mode: str = "auto",
773
774
        trust_remote_code: bool = True,
        seed: Optional[int] = 0,
775
        max_model_len: int = 1024,
776
        dtype: str = "auto",
777
        disable_log_stats: bool = True,
778
        tensor_parallel_size: int = 1,
779
        block_size: int = 16,
780
        enable_chunked_prefill: Optional[bool] = False,
781
        swap_space: int = 4,
782
        enforce_eager: Optional[bool] = False,
783
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
784
785
786
    ) -> None:
        self.model = LLM(
            model=model_name,
787
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
788
            tokenizer=tokenizer_name,
789
            tokenizer_mode=tokenizer_mode,
790
            trust_remote_code=trust_remote_code,
Woosuk Kwon's avatar
Woosuk Kwon committed
791
            dtype=dtype,
792
            seed=seed,
793
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
794
            enforce_eager=enforce_eager,
795
            disable_log_stats=disable_log_stats,
796
            tensor_parallel_size=tensor_parallel_size,
797
            max_model_len=max_model_len,
798
799
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
800
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
801
802
        )

803
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
804
        self,
805
        prompts: Union[list[str], list[torch.Tensor]],
806
        images: Optional[PromptImageInput] = None,
807
808
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
809
    ) -> list[TextPrompt]:
810

811
812
813
814
815
        if any(x is not None and len(x) != len(prompts)
               for x in [images, videos, audios]):
            raise ValueError(
                "All non-None multimodal inputs must have the same length as "
                "prompts")
816

817
818
819
820
821
822
823
824
825
826
        inputs = []
        for i, prompt in enumerate(prompts):
            multi_modal_data = {}
            if images is not None and (image := images[i]) is not None:
                multi_modal_data["image"] = image
            if videos is not None and (video := videos[i]) is not None:
                multi_modal_data["video"] = video
            if audios is not None and (audio := audios[i]) is not None:
                multi_modal_data["audio"] = audio

827
828
829
830
831
832
            text_prompt_kwargs = {
                ("prompt" if isinstance(prompt, str) else "prompt_embeds"):
                prompt,
                "multi_modal_data": multi_modal_data or None
            }
            inputs.append(TextPrompt(**text_prompt_kwargs))
833
834
835
836
837

        return inputs

    def generate(
        self,
838
        prompts: Union[list[str], list[torch.Tensor]],
839
840
841
842
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
843
        **kwargs: Any,
844
    ) -> list[tuple[list[list[int]], list[str]]]:
845
846
847
848
849
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

850
        req_outputs = self.model.generate(inputs,
851
852
                                          sampling_params=sampling_params,
                                          **kwargs)
853

854
        outputs: list[tuple[list[list[int]], list[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
855
856
857
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
858
859
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
860
861
            for sample in req_output.outputs:
                output_str = sample.text
862
                output_ids = list(sample.token_ids)
863
                req_sample_output_ids.append(prompt_ids + output_ids)
864
                req_sample_output_strs.append((prompt_str or "") + output_str)
865
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
866
867
        return outputs

868
    @staticmethod
869
    def _final_steps_generate_w_logprobs(
870
871
872
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
873
        for req_output in req_outputs:
874
            assert len(req_output.outputs) > 0
875
876
            for sample in req_output.outputs:
                output_str = sample.text
877
                output_ids = list(sample.token_ids)
878
                output_logprobs = sample.logprobs
879
880
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
881
882
        return outputs

883
884
    def generate_w_logprobs(
        self,
885
        prompts: list[str],
886
        sampling_params: SamplingParams,
887
888
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
889
        videos: Optional[PromptVideoInput] = None,
890
        **kwargs: Any,
891
892
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
893
894
895
896
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
897

898
        req_outputs = self.model.generate(inputs,
899
900
                                          sampling_params=sampling_params,
                                          **kwargs)
901
902
903
904
905
906
907

        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        return ([x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
                if sampling_params.prompt_logprobs is None else
                toks_str_logsprobs_prompt_logprobs)
908
909
910

    def generate_encoder_decoder_w_logprobs(
        self,
911
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
912
        sampling_params: SamplingParams,
913
914
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
915
916
917
918
919
920
921
        '''
        Logprobs generation for vLLM encoder/decoder models
        '''

        assert sampling_params.logprobs is not None
        req_outputs = self.model.generate(encoder_decoder_prompts,
                                          sampling_params=sampling_params)
922
923
924
925
926
927
        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        return ([x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
                if sampling_params.prompt_logprobs is None else
                toks_str_logsprobs_prompt_logprobs)
928

Woosuk Kwon's avatar
Woosuk Kwon committed
929
930
    def generate_greedy(
        self,
931
        prompts: Union[list[str], list[torch.Tensor]],
Woosuk Kwon's avatar
Woosuk Kwon committed
932
        max_tokens: int,
933
        images: Optional[PromptImageInput] = None,
934
935
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
936
        **kwargs: Any,
937
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
938
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
939
940
941
942
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
943
944
                                audios=audios,
                                **kwargs)
945
946
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
947

948
949
    def generate_greedy_logprobs(
        self,
950
        prompts: list[str],
951
952
        max_tokens: int,
        num_logprobs: int,
953
        num_prompt_logprobs: Optional[int] = None,
954
955
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
956
        videos: Optional[PromptVideoInput] = None,
957
958
        stop_token_ids: Optional[list[int]] = None,
        stop: Optional[list[str]] = None,
959
        **kwargs: Any,
960
961
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
962
963
964
965
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
966
            prompt_logprobs=num_prompt_logprobs,
967
968
            stop_token_ids=stop_token_ids,
            stop=stop)
969
970
971
972
973

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
974
975
                                        videos=videos,
                                        **kwargs)
976

977
978
    def generate_encoder_decoder_greedy_logprobs(
        self,
979
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
980
981
        max_tokens: int,
        num_logprobs: int,
982
        num_prompt_logprobs: Optional[int] = None,
983
        skip_special_tokens: bool = True,
984
985
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
986
987
988
989
990
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
991
            skip_special_tokens=skip_special_tokens,
992
        )
993
994
995
996
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

997
        return self.generate_encoder_decoder_w_logprobs(
998
999
            encoder_decoder_prompts, greedy_logprobs_params)

1000
    def generate_beam_search(
1001
        self,
1002
        prompts: list[str],
1003
1004
        beam_width: int,
        max_tokens: int,
1005
1006
1007
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
1008
    ) -> list[tuple[list[list[int]], list[str]]]:
1009
1010
1011
1012
1013
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

1014
        outputs = self.model.beam_search(
1015
            inputs,
1016
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
1017
1018
1019
1020
1021
1022
1023
        returned_outputs = []
        for output in outputs:
            token_ids = [x.tokens for x in output.sequences]
            texts = [x.text for x in output.sequences]
            returned_outputs.append((token_ids, texts))
        return returned_outputs

1024
    def classify(self, prompts: list[str]) -> list[list[float]]:
1025
1026
1027
        req_outputs = self.model.classify(prompts)
        return [req_output.outputs.probs for req_output in req_outputs]

1028
1029
1030
1031
1032
1033
1034
    def encode(self,
               prompts: list[str],
               images: Optional[PromptImageInput] = None,
               videos: Optional[PromptVideoInput] = None,
               audios: Optional[PromptAudioInput] = None,
               *args,
               **kwargs) -> list[list[float]]:
Cyrus Leung's avatar
Cyrus Leung committed
1035
1036
1037
1038
1039
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

1040
        req_outputs = self.model.embed(inputs, *args, **kwargs)
Cyrus Leung's avatar
Cyrus Leung committed
1041
        return [req_output.outputs.embedding for req_output in req_outputs]
1042

1043
1044
    def score(
        self,
1045
1046
        text_1: Union[str, list[str]],
        text_2: Union[str, list[str]],
1047
1048
        *args,
        **kwargs,
1049
    ) -> list[float]:
1050
        req_outputs = self.model.score(text_1, text_2, *args, **kwargs)
1051
        return [req_output.outputs.score for req_output in req_outputs]
1052

1053
1054
1055
1056
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
        executor = self.model.llm_engine.model_executor
        return executor.apply_model(func)

1057
1058
1059
1060
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1061
        del self.model
1062
        cleanup_dist_env_and_memory()
1063

Woosuk Kwon's avatar
Woosuk Kwon committed
1064

1065
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1066
1067
def vllm_runner():
    return VllmRunner
1068
1069


1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
    logger = logging.getLogger("vllm")
    logger.propagate = True
    yield
    logger.propagate = False


@pytest.fixture()
def caplog_vllm(temporary_enable_log_propagate, caplog):
    # To capture vllm log, we should enable propagate=True temporarily
    # because caplog depends on logs propagated to the root logger.
    yield caplog
1084
1085
1086
1087
1088
1089
1090


@pytest.fixture(scope="session")
def num_gpus_available():
    """Get number of GPUs without initializing the CUDA context
    in current process."""

1091
    return cuda_device_count_stateless()
1092
1093
1094


temp_dir = tempfile.gettempdir()
1095
1096
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1097
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1098
1099
1100
1101


@pytest.fixture
def dummy_opt_path():
1102
1103
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1104
        snapshot_download(repo_id="facebook/opt-125m",
1105
                          local_dir=_dummy_opt_path,
1106
1107
1108
1109
1110
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1111
        with open(json_path) as f:
1112
1113
1114
1115
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
    return _dummy_opt_path


@pytest.fixture
def dummy_llava_path():
    json_path = os.path.join(_dummy_llava_path, "config.json")
    if not os.path.exists(_dummy_llava_path):
        snapshot_download(repo_id="llava-hf/llava-1.5-7b-hf",
                          local_dir=_dummy_llava_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1130
        with open(json_path) as f:
1131
1132
1133
1134
1135
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148


@pytest.fixture
def dummy_gemma2_embedding_path():
    json_path = os.path.join(_dummy_gemma2_embedding_path, "config.json")
    if not os.path.exists(_dummy_gemma2_embedding_path):
        snapshot_download(repo_id="BAAI/bge-multilingual-gemma2",
                          local_dir=_dummy_gemma2_embedding_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1149
        with open(json_path) as f:
1150
1151
1152
1153
1154
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
    parser.addoption("--optional",
                     action="store_true",
                     default=False,
                     help="run optional test")


def pytest_collection_modifyitems(config, items):
    if config.getoption("--optional"):
        # --optional given in cli: do not skip optional tests
        return
    skip_optional = pytest.mark.skip(reason="need --optional option to run")
    for item in items:
        if "optional" in item.keywords:
1173
            item.add_marker(skip_optional)
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185


@pytest.fixture(scope="session")
def cli_config_file():
    """Return the path to the CLI config file."""
    return os.path.join(_TEST_DIR, "config", "test_config.yaml")


@pytest.fixture(scope="session")
def cli_config_file_with_model():
    """Return the path to the CLI config file with model."""
    return os.path.join(_TEST_DIR, "config", "test_config_with_model.yaml")