conftest.py 41.2 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
zhuwenwen's avatar
zhuwenwen committed
7

8
from typing import Any, Callable, Optional, TypedDict, TypeVar, Union
9
10
import pytest
import pytest_html
Woosuk Kwon's avatar
Woosuk Kwon committed
11

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

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

40
from vllm.utils import cuda_device_count_stateless
zhuwenwen's avatar
zhuwenwen committed
41
from .utils import models_path_prefix
42
43


44
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
45

46
47
48
_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")]
49
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
50

Cyrus Leung's avatar
Cyrus Leung committed
51
_M = TypeVar("_M")
52

53
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
54
55

PromptImageInput = _PromptMultiModalInput[Image.Image]
56
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
57
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
58

59

60
def _read_prompts(filename: str) -> list[str]:
61
    with open(filename) as f:
62
63
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
64
65


66
class ImageAssetPrompts(TypedDict):
67
68
    stop_sign: str
    cherry_blossom: str
69
70


71
class ImageTestAssets(list[ImageAsset]):
72
73

    def __init__(self) -> None:
74
75
76
77
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
78

79
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
80
81
82
83
84
85
        """
        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.
        """
86
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
87
88


89
90
class VideoAssetPrompts(TypedDict):
    baby_reading: str
91
92


93
class VideoTestAssets(list[VideoAsset]):
94
95
96

    def __init__(self) -> None:
        super().__init__([
97
            VideoAsset("baby_reading"),
98
99
        ])

100
101
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
102
103


104
class AudioAssetPrompts(TypedDict):
105
106
    mary_had_lamb: str
    winning_call: str
107
108


109
class AudioTestAssets(list[AudioAsset]):
110
111
112
113
114
115
116

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

117
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
118
        return [prompts["mary_had_lamb"], prompts["winning_call"]]
119
120


121
IMAGE_ASSETS = ImageTestAssets()
122
"""Singleton instance of {class}`ImageTestAssets`."""
123
VIDEO_ASSETS = VideoTestAssets()
124
"""Singleton instance of {class}`VideoTestAssets`."""
125
AUDIO_ASSETS = AudioTestAssets()
126
"""Singleton instance of {class}`AudioTestAssets`."""
127
128


129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
@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
149
@pytest.fixture(params=[True, False])
150
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
151
152
153
154
155
156
157
158
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
159
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
160
    else:
161
162
163
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
164
165


166
167
168
169
170
171
172
@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


173
174
175
176
177
178
179
180
181
182
183
184
@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
185
    cleanup_dist_env_and_memory()
186
187


188
@pytest.fixture()
189
def should_do_global_cleanup_after_test(request) -> bool:
190
191
192
193
    """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.
    """
194

195
    return not request.node.get_closest_marker("skip_global_cleanup")
196
197


198
@pytest.fixture(autouse=True)
199
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
200
    yield
201
    if should_do_global_cleanup_after_test:
202
        cleanup_dist_env_and_memory()
203
204


205
206
207
208
209
210
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


219
220
221
222
223
224
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


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


232
@pytest.fixture
233
def example_encoder_decoder_prompts(
234
) -> dict[DecoderPromptType, list[ExplicitEncoderDecoderPrompt]]:
235
236
237
238
239
240
    '''
    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:
241

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    * 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:
257
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
258
        DecoderPromptType.EMPTY_STR:
259
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
260
        DecoderPromptType.CUSTOM:
261
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
262
263
264
    }


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


273
@pytest.fixture(scope="session")
274
def image_assets() -> ImageTestAssets:
275
276
277
    return IMAGE_ASSETS


278
@pytest.fixture(scope="session")
279
def video_assets() -> VideoTestAssets:
280
281
282
    return VIDEO_ASSETS


283
@pytest.fixture(scope="session")
284
def audio_assets() -> AudioTestAssets:
285
286
287
    return AUDIO_ASSETS


288
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
289
_R = TypeVar("_R")
290

Woosuk Kwon's avatar
Woosuk Kwon committed
291
292
293

class HfRunner:

294
    def get_default_device(self):
295
        from vllm.platforms import current_platform
296

297
298
        return ("cpu"
                if current_platform.is_cpu() else current_platform.device_type)
299
300

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

304
        if device is None:
305
            device = self.device
306

307
308
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
309

310
311
312
313
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
314

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

329
330
        self.config = AutoConfig.from_pretrained(
            model_name,
331
            trust_remote_code=trust_remote_code,
332
333
        )
        self.device = self.get_default_device()
334
335
336
337
338
339
        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,
        )
340
341
342
343

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

344
        if is_sentence_transformer:
345
346
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
347
348
349
350
351

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

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

371
372
373
374
375
376
            # 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)

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

            self.model = model
383

384
385
386
387
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
388
                trust_remote_code=trust_remote_code,
389
            )
390

391
392
393
394
395
396
        # 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,
397
            trust_remote_code=trust_remote_code,
398
        )
399
400
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
401

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

412
413
414
415
416
417
        if videos is not None:
            assert len(prompts) == len(videos)

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

418
        all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
419
        for i, prompt in enumerate(prompts):
420
            processor_kwargs: dict[str, Any] = {
421
422
423
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
424
425
426
427
            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
428
429
430
431
432
433
434
435
436
            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
437
438

            inputs = self.processor(**processor_kwargs)
439
440
            if isinstance(inputs, BatchFeature):
                inputs = inputs.to(dtype=self.dtype)
441

442
443
444
445
            all_inputs.append(inputs)

        return all_inputs

446
447
448
449
450
451
452
453
454
    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

455
    def classify(self, prompts: list[str]) -> list[str]:
456
457
458
459
460
461
462
463
464
465
        # 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

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

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

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

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
514
515
516

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

533
534
535
536
537
538
539
540
541
        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
542

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

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

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

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

594
595
596
597
        return seq_logprobs

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

604
        # convert to dict
605
        seq_logprobs_lst: list[dict[int, float]] = []
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
        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,
        )

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

638
639
640
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
641

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

653
654
655
656
657
658
659
660
661
662
663
664
            (
                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))
665

666
667
668
669
670
671
        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,
672
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
673
674
        max_tokens: int,
        num_logprobs: int,
675
        images: Optional[PromptImageInput] = None,
676
        **kwargs: Any,
677
    ) -> list[TokensTextLogprobs]:
678
679
680
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
681

682
683
684
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
685

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

695
696
            encoder_inputs = self.processor(**processor_kwargs)
            encoder_inputs = self.wrap_device(encoder_inputs)
697
698
699
700

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

            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,
712
                **encoder_inputs,
713
714
715
716
717
718
719
720
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
721
722
723
724
725
726
727
728
729
730
731

            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]

732
733
734
    def encode(self, prompts: list[str], *args,
               **kwargs) -> list[list[torch.Tensor]]:
        return self.model.encode(prompts, *args, **kwargs)
735

736
    def predict(self, prompts: list[list[str]]) -> torch.Tensor:
737
738
        return self.model.predict(prompts, convert_to_tensor=True)

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

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

Woosuk Kwon's avatar
Woosuk Kwon committed
746

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


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

757
758
759
760
761
762
    - `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.
763
    - `enforce_eager`: Set to `False` to test CUDA graph.
764
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
765
766
767
768

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

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

810
811
812
813
814
        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")
815

816
817
818
819
820
821
822
823
824
        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
825

826
827
828
829
830
831
            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))
832
833
834
835
836

        return inputs

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

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

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

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

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

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

        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)
907
908
909

    def generate_encoder_decoder_w_logprobs(
        self,
910
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
911
        sampling_params: SamplingParams,
912
913
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
914
915
916
917
918
919
920
        '''
        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)
921
922
923
924
925
926
        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)
927

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

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

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

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

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

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

1013
        outputs = self.model.beam_search(
1014
            inputs,
1015
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
1016
1017
1018
1019
1020
1021
1022
        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

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

1027
1028
1029
1030
1031
1032
1033
    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
1034
1035
1036
1037
1038
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

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

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

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

1054
1055
1056
1057
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1058
        del self.model
1059
        cleanup_dist_env_and_memory()
1060

Woosuk Kwon's avatar
Woosuk Kwon committed
1061

1062
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1063
1064
def vllm_runner():
    return VllmRunner
1065
1066


1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
@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
1081
1082
1083
1084
1085
1086
1087


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

1088
    return cuda_device_count_stateless()
1089
1090


1091
# temp_dir = tempfile.gettempdir()
zhuwenwen's avatar
zhuwenwen committed
1092
1093
1094
1095
_dummy_opt_path = os.path.join(models_path_prefix, "dummy_opt")
_dummy_llava_path = os.path.join(models_path_prefix, "dummy_llava")
_dummy_gemma2_embedding_path = os.path.join(models_path_prefix, "dummy_gemma2_embedding")

1096
1097
1098
1099


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

1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133

# 定义一个 pytest 钩子,在测试后生成报告
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    # 获取测试结果
    outcome = yield
    result = outcome.get_result()

    # 如果测试失败并且有浏览器实例,添加截图
    if result.when == "call" and result.failed:
        if hasattr(item, "funcargs") and "browser" in item.funcargs:
            browser = item.funcargs["browser"]
            screenshot_path = "screenshot.png"  # 设置截图路径
            browser.save_screenshot(screenshot_path)

            # 如果测试结果有 extra 属性,则添加截图
            if hasattr(result, "extra"):
                result.extra.append(pytest_html.extras.image(screenshot_path))
zhuwenwen's avatar
zhuwenwen committed
1134
1135


1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
@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)
1147
        with open(json_path) as f:
1148
1149
1150
1151
1152
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165


@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)
1166
        with open(json_path) as f:
1167
1168
1169
1170
1171
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190


# 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:
            item.add_marker(skip_optional)
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202


@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")