conftest.py 36.5 KB
Newer Older
1
import json
2
import os
3
import tempfile
4
from collections import UserList
5
from enum import Enum
6
7
from typing import (Any, Callable, Dict, List, Optional, Tuple, Type,
                    TypedDict, TypeVar, Union)
8
9
import pytest
import pytest_html
Woosuk Kwon's avatar
Woosuk Kwon committed
10

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

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

42
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
43

44
45
46
_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")]
47

Cyrus Leung's avatar
Cyrus Leung committed
48
49
50
51
52
53
_M = TypeVar("_M")
_PromptMultiModalInput = Union[List[_M], List[List[_M]]]

PromptImageInput = _PromptMultiModalInput[Image.Image]
PromptAudioInput = _PromptMultiModalInput[Tuple[np.ndarray, int]]
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
54

55

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


62
63
64
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
65
66


67
68
class _ImageAssetsBase(UserList[ImageAsset]):
    pass
69

70
71

class _ImageAssets(_ImageAssetsBase):
72
73

    def __init__(self) -> None:
74
75
76
77
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
78
79
80
81
82
83
84
85

    def prompts(self, prompts: _ImageAssetPrompts) -> List[str]:
        """
        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
91
92
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


93
94
class _VideoAssetsBase(UserList[VideoAsset]):
    pass
95
96
97
98
99
100
101
102
103
104
105
106
107


class _VideoAssets(_VideoAssetsBase):

    def __init__(self) -> None:
        super().__init__([
            VideoAsset("sample_demo_1.mp4"),
        ])

    def prompts(self, prompts: _VideoAssetPrompts) -> List[str]:
        return [prompts["sample_demo_1"]]


108
109
IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""
110
111
VIDEO_ASSETS = _VideoAssets()
"""Singleton instance of :class:`_VideoAssets`."""
112
113


Joe Runde's avatar
Joe Runde committed
114
@pytest.fixture(params=[True, False])
115
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
116
117
118
119
120
121
122
123
    # 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")
124
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
125
    else:
126
127
128
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
129
130


131
132
133
134
135
136
137
@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


138
139
140
141
142
143
144
145
146
147
148
149
@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
150
    cleanup_dist_env_and_memory()
151
152


153
@pytest.fixture()
154
def should_do_global_cleanup_after_test(request) -> bool:
155
156
157
158
    """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.
    """
159

160
    return not request.node.get_closest_marker("skip_global_cleanup")
161
162


163
@pytest.fixture(autouse=True)
164
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
165
    yield
166
    if should_do_global_cleanup_after_test:
167
        cleanup_dist_env_and_memory()
168
169


170
171
172
173
174
175
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
176
177
@pytest.fixture
def example_prompts() -> List[str]:
178
179
    prompts = []
    for filename in _TEST_PROMPTS:
180
        prompts += _read_prompts(filename)
181
182
183
    return prompts


184
185
186
187
188
189
190
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


191
@pytest.fixture
192
193
def example_encoder_decoder_prompts(
) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
194
195
196
197
198
199
    '''
    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:
200

201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    * 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:
216
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
217
        DecoderPromptType.EMPTY_STR:
218
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
219
        DecoderPromptType.CUSTOM:
220
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
221
222
223
    }


224
225
226
227
@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
228
        prompts += _read_prompts(filename)
229
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
230
231


232
233
234
235
236
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


237
238
239
240
241
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


242
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
243

Woosuk Kwon's avatar
Woosuk Kwon committed
244
245
246

class HfRunner:

247
    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
248
249
250
        if x is None or isinstance(x, (bool, )):
            return x

251
        if device is None:
252
            device = "cpu" if current_platform.is_cpu() else "cuda"
253

254
255
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
256

257
258
259
260
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
261

Woosuk Kwon's avatar
Woosuk Kwon committed
262
263
264
265
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
266
        *,
267
        model_kwargs: Optional[Dict[str, Any]] = None,
268
        is_sentence_transformer: bool = False,
269
        is_cross_encoder: bool = False,
270
        skip_tokenizer_init: bool = False,
271
        auto_cls: Type[_BaseAutoModelClass] = AutoModelForCausalLM,
272
        postprocess_inputs: Callable[..., BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
273
    ) -> None:
274
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
275

276
        self.model_name = model_name
277

278
        if is_sentence_transformer:
279
280
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
281
282
283
284
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
285
                    trust_remote_code=True,
286
                ).to(dtype=torch_dtype))
287
288
289
290
291
292
293
294
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
            self.model = CrossEncoder(model_name,
                                      device="cpu",
                                      trust_remote_code=True)
            self.model.model = self.wrap_device(self.model.model)\
                .to(dtype=torch_dtype)
295
        else:
296
            model_kwargs = model_kwargs if model_kwargs is not None else {}
297
            self.model = self.wrap_device(
298
                auto_cls.from_pretrained(
299
300
301
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
302
                    **model_kwargs,
303
                ))
304

305
306
307
308
309
310
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
311

312
313
314
315
316
317
318
319
        # 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,
            trust_remote_code=True,
        )
320
321
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
322

323
        self.dtype = dtype
324
325
        self.postprocess_inputs = postprocess_inputs

326
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
327
328
        self,
        prompts: List[str],
329
        images: Optional[PromptImageInput] = None,
330
331
332
333
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[BatchEncoding]:
        if images is not None:
334
            assert len(prompts) == len(images)
335

336
337
338
339
340
341
342
        if videos is not None:
            assert len(prompts) == len(videos)

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

        all_inputs: List[BatchEncoding] = []
343
        for i, prompt in enumerate(prompts):
344
345
346
347
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
348
349
350
351
352
353
            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
            if audios is not None and (audio_tuple := audios[i]) is not None:
                audio, sr = audio_tuple
354
355
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
356
357

            inputs = self.processor(**processor_kwargs)
358
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
359

360
361
362
363
            all_inputs.append(inputs)

        return all_inputs

364
365
366
367
368
369
370
371
372
373
374
    def classify(self, prompts: List[str]) -> List[str]:
        # 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

375
376
377
378
    def generate(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
379
        videos: Optional[PromptVideoInput] = None,
380
381
382
383
384
385
386
        audios: Optional[PromptAudioInput] = None,
        **kwargs: Any,
    ) -> List[Tuple[List[List[int]], List[str]]]:
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
387

388
389
        outputs: List[Tuple[List[List[int]], List[str]]] = []
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
390
            output_ids = self.model.generate(
391
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
392
393
394
                use_cache=True,
                **kwargs,
            )
395
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
396
397
398
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
399
400
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
401
402
403
404
405
406
407
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
408
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
409
        videos: Optional[PromptVideoInput] = None,
410
        audios: Optional[PromptAudioInput] = None,
411
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
412
    ) -> List[Tuple[List[int], str]]:
413
414
        outputs = self.generate(prompts,
                                do_sample=False,
415
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
416
                                images=images,
417
418
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
419
                                **kwargs)
420
421
422

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
423
424
425
426
427
428

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
429
    ) -> List[Tuple[List[List[int]], List[str]]]:
430
431
432
433
434
435
436
437
438
439
440
441
442
443
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
                                num_return_sequences=beam_width)
        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
444

445
446
447
448
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
449
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
450
        videos: Optional[PromptVideoInput] = None,
451
        audios: Optional[PromptAudioInput] = None,
452
        **kwargs: Any,
453
    ) -> List[List[torch.Tensor]]:
454
455
456
457
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
458

459
460
        all_logprobs: List[List[torch.Tensor]] = []
        for inputs in all_inputs:
461
            output = self.model.generate(
462
                **self.wrap_device(inputs, device=self.model.device.type),
463
464
465
466
467
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
468
                **kwargs,
469
            )
470
471
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
472
473
474
            all_logprobs.append(seq_logprobs)
        return all_logprobs

475
    def _hidden_states_to_seq_logprobs(
476
        self,
477
478
479
480
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

481
482
483
484
        seq_logprobs: List[torch.Tensor] = []
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
485
486
                last_hidden_states.to(output_embeddings.weight.device),
                output_embeddings.weight.t(),
487
            )
488
489
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
490
491
492
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

493
494
495
496
497
498
499
500
501
502
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
        num_logprobs: int,
    ) -> Tuple[List[Dict[int, float]], int]:
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
        # convert to dict
        seq_logprobs_lst: List[Dict[int, float]] = []
        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,
        )

522
523
524
525
526
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
527
528
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
529
        videos: Optional[PromptVideoInput] = None,
530
        **kwargs: Any,
531
    ) -> List[TokensTextLogprobs]:
532
533
534
535
536
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

537
538
539
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
540

541
        for inputs in all_inputs:
542
            output = self.model.generate(
543
                **self.wrap_device(inputs, device=self.model.device.type),
544
545
546
547
548
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
549
                **kwargs,
550
551
            )

552
553
554
555
556
557
558
559
560
561
562
563
            (
                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))
564

565
566
567
568
569
570
        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,
571
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
572
573
        max_tokens: int,
        num_logprobs: int,
574
        images: Optional[PromptImageInput] = None,
575
        **kwargs: Any,
576
    ) -> List[TokensTextLogprobs]:
577
578
579
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
580

581
582
583
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
584

585
586
587
588
589
590
591
592
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
            processor_kwargs: Dict[str, Any] = {
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
593

594
            encoder_input_ids = self.wrap_device(
595
                self.processor(**processor_kwargs).input_ids,
596
597
598
599
600
601
602
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
603
                    self.tokenizer(decoder_prompt,
604
605
606
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623

            output = self.model.generate(
                encoder_input_ids,
                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,
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
624
625
626
627
628
629
630
631
632
633
634

            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]

635
636
637
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

638
639
640
    def predict(self, prompts: List[List[str]]) -> torch.Tensor:
        return self.model.predict(prompts, convert_to_tensor=True)

641
642
643
644
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
645
        del self.model
646
        cleanup_dist_env_and_memory()
647

Woosuk Kwon's avatar
Woosuk Kwon committed
648

Cyrus Leung's avatar
Cyrus Leung committed
649
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
650
651
652
653
654
655
656
657
658
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
659
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
660
        tokenizer_name: Optional[str] = None,
661
        tokenizer_mode: str = "auto",
662
663
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
664
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
665
        dtype: str = "half",
666
        disable_log_stats: bool = True,
667
        tensor_parallel_size: int = 1,
668
669
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
670
        swap_space: int = 4,
671
        enforce_eager: Optional[bool] = False,
672
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
673
674
675
    ) -> None:
        self.model = LLM(
            model=model_name,
676
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
677
            tokenizer=tokenizer_name,
678
            tokenizer_mode=tokenizer_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
679
680
            trust_remote_code=True,
            dtype=dtype,
681
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
682
            enforce_eager=enforce_eager,
683
            disable_log_stats=disable_log_stats,
684
            tensor_parallel_size=tensor_parallel_size,
685
            max_model_len=max_model_len,
686
687
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
688
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
689
690
        )

691
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
692
693
        self,
        prompts: List[str],
694
        images: Optional[PromptImageInput] = None,
695
696
697
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[TextPrompt]:
698
        if images is not None:
699
            assert len(prompts) == len(images)
700

701
702
703
704
705
706
        if videos is not None:
            assert len(prompts) == len(videos)

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

707
708
709
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
Cyrus Leung's avatar
Cyrus Leung committed
710
711
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
712

713
714
        if videos is not None:
            for i, video in enumerate(videos):
Cyrus Leung's avatar
Cyrus Leung committed
715
716
                if video is not None:
                    inputs[i]["multi_modal_data"] = {"video": video}
717
718
719

        if audios is not None:
            for i, audio in enumerate(audios):
Cyrus Leung's avatar
Cyrus Leung committed
720
721
                if audio is not None:
                    inputs[i]["multi_modal_data"] = {"audio": audio}
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736

        return inputs

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[Tuple[List[List[int]], List[str]]]:
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
737

738
        req_outputs = self.model.generate(inputs,
739
                                          sampling_params=sampling_params)
740
741

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
742
743
744
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
745
746
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
747
748
            for sample in req_output.outputs:
                output_str = sample.text
749
                output_ids = list(sample.token_ids)
750
751
752
                req_sample_output_ids.append(prompt_ids + output_ids)
                req_sample_output_strs.append(prompt_str + output_str)
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
753
754
        return outputs

755
    @staticmethod
756
757
    def _final_steps_generate_w_logprobs(
        req_outputs: List[RequestOutput],
758
759
    ) -> List[TokensTextLogprobsPromptLogprobs]:
        outputs: List[TokensTextLogprobsPromptLogprobs] = []
760
        for req_output in req_outputs:
761
            assert len(req_output.outputs) > 0
762
763
            for sample in req_output.outputs:
                output_str = sample.text
764
                output_ids = list(sample.token_ids)
765
                output_logprobs = sample.logprobs
766
767
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
768
769
        return outputs

770
771
772
773
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
774
775
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
776
        videos: Optional[PromptVideoInput] = None,
777
778
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
779
780
781
782
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
783

784
        req_outputs = self.model.generate(inputs,
785
                                          sampling_params=sampling_params)
786
787
788
789
790
791
792

        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)
793
794
795

    def generate_encoder_decoder_w_logprobs(
        self,
796
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
797
        sampling_params: SamplingParams,
798
799
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
800
801
802
803
804
805
806
        '''
        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)
807
808
809
810
811
812
        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)
813

Woosuk Kwon's avatar
Woosuk Kwon committed
814
815
816
817
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
818
        images: Optional[PromptImageInput] = None,
819
820
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
821
822
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
823
824
825
826
827
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
                                audios=audios)
828
829
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
830

831
832
833
834
835
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
836
        num_prompt_logprobs: Optional[int] = None,
837
838
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
839
        videos: Optional[PromptVideoInput] = None,
840
        stop_token_ids: Optional[List[int]] = None,
841
        stop: Optional[List[str]] = None,
842
843
844
845
846
847
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
848
            prompt_logprobs=num_prompt_logprobs,
849
850
            stop_token_ids=stop_token_ids,
            stop=stop)
851
852
853
854
855
856

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
                                        videos=videos)
857

858
859
    def generate_encoder_decoder_greedy_logprobs(
        self,
860
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
861
862
        max_tokens: int,
        num_logprobs: int,
863
864
865
866
867
868
869
870
871
        num_prompt_logprobs: Optional[int] = None,
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
        )
872
873
874
875
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

876
        return self.generate_encoder_decoder_w_logprobs(
877
878
            encoder_decoder_prompts, greedy_logprobs_params)

879
    def generate_beam_search(
880
881
882
883
884
        self,
        prompts: Union[List[str], List[List[int]]],
        beam_width: int,
        max_tokens: int,
    ) -> List[Tuple[List[List[int]], List[str]]]:
885
886
887
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
888
889
890
891
892
893
894
        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

895
896
897
898
    def classify(self, prompts: List[str]) -> List[List[float]]:
        req_outputs = self.model.classify(prompts)
        return [req_output.outputs.probs for req_output in req_outputs]

Cyrus Leung's avatar
Cyrus Leung committed
899
900
901
902
903
904
905
906
907
908
909
910
    def encode(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[List[float]]:
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

911
        req_outputs = self.model.embed(inputs)
Cyrus Leung's avatar
Cyrus Leung committed
912
        return [req_output.outputs.embedding for req_output in req_outputs]
913

914
915
916
917
    def score(
        self,
        text_1: Union[str, List[str]],
        text_2: Union[str, List[str]],
918
    ) -> List[float]:
919
        req_outputs = self.model.score(text_1, text_2)
920
        return [req_output.outputs.score for req_output in req_outputs]
921

922
923
924
925
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
926
        del self.model
927
        cleanup_dist_env_and_memory()
928

Woosuk Kwon's avatar
Woosuk Kwon committed
929

930
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
931
932
def vllm_runner():
    return VllmRunner
933
934
935
936
937
938
939
940
941


def get_tokenizer_pool_config(tokenizer_group_type):
    if tokenizer_group_type is None:
        return None
    if tokenizer_group_type == "ray":
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type="ray",
                                   extra_config={})
942
943
944
945
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
946
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962


@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
963
964
965
966
967
968
969


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

970
    return cuda_device_count_stateless()
971
972


973
# temp_dir = tempfile.gettempdir()
zhuwenwen's avatar
zhuwenwen committed
974
975
976
977
_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")

978
979
980
981


@pytest.fixture
def dummy_opt_path():
982
983
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
984
        snapshot_download(repo_id="facebook/opt-125m",
985
                          local_dir=_dummy_opt_path,
986
987
988
989
990
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
991
        with open(json_path) as f:
992
993
994
995
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
996
997
    return _dummy_opt_path

998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015

# 定义一个 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
1016
1017


1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
@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)
1029
        with open(json_path) as f:
1030
1031
1032
1033
1034
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047


@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)
1048
        with open(json_path) as f:
1049
1050
1051
1052
1053
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072


# 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)