conftest.py 35.8 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)
Woosuk Kwon's avatar
Woosuk Kwon committed
8

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

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

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

41
42
43
_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")]
44

Cyrus Leung's avatar
Cyrus Leung committed
45
46
47
48
49
50
_M = TypeVar("_M")
_PromptMultiModalInput = Union[List[_M], List[List[_M]]]

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

52

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


59
60
61
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
62
63


64
65
class _ImageAssetsBase(UserList[ImageAsset]):
    pass
66

67
68

class _ImageAssets(_ImageAssetsBase):
69
70

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

    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.
        """
83
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
84
85


86
87
88
89
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


90
91
class _VideoAssetsBase(UserList[VideoAsset]):
    pass
92
93
94
95
96
97
98
99
100
101
102
103
104


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"]]


105
106
IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""
107
108
VIDEO_ASSETS = _VideoAssets()
"""Singleton instance of :class:`_VideoAssets`."""
109
110


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

    yield
Joe Runde's avatar
Joe Runde committed
126
127


128
129
130
131
132
133
134
@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


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


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

157
    return not request.node.get_closest_marker("skip_global_cleanup")
158
159


160
@pytest.fixture(autouse=True)
161
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
162
    yield
163
    if should_do_global_cleanup_after_test:
164
        cleanup_dist_env_and_memory()
165
166


167
168
169
170
171
172
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


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


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

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


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


229
230
231
232
233
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


234
235
236
237
238
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


239
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
240

Woosuk Kwon's avatar
Woosuk Kwon committed
241
242
243

class HfRunner:

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

248
        if device is None:
249
            device = "cpu" if current_platform.is_cpu() else "cuda"
250

251
252
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
253

254
255
256
257
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
258

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

273
        self.model_name = model_name
274

275
        if is_sentence_transformer:
276
277
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
278
279
280
281
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
282
                    trust_remote_code=True,
283
                ).to(dtype=torch_dtype))
284
285
286
287
288
289
290
291
        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)
292
        else:
293
            model_kwargs = model_kwargs if model_kwargs is not None else {}
294
            self.model = self.wrap_device(
295
                auto_cls.from_pretrained(
296
297
298
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
299
                    **model_kwargs,
300
                ))
301

302
303
304
305
306
307
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
308

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

320
        self.dtype = dtype
321
322
        self.postprocess_inputs = postprocess_inputs

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

333
334
335
336
337
338
339
        if videos is not None:
            assert len(prompts) == len(videos)

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

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

            inputs = self.processor(**processor_kwargs)
355
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
356

357
358
359
360
            all_inputs.append(inputs)

        return all_inputs

361
362
363
364
365
366
367
368
369
370
371
    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

372
373
374
375
    def generate(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
376
        videos: Optional[PromptVideoInput] = None,
377
378
379
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)

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

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

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
420
421
422
423
424
425

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
426
    ) -> List[Tuple[List[List[int]], List[str]]]:
427
428
429
430
431
432
433
434
435
436
437
438
439
440
        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
441

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

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

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

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

490
491
492
493
494
495
496
497
498
499
        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)

500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
        # 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,
        )

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

534
535
536
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
537

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

549
550
551
552
553
554
555
556
557
558
559
560
            (
                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))
561

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

578
579
580
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
581

582
583
584
585
586
587
588
589
        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]
590

591
            encoder_input_ids = self.wrap_device(
592
                self.processor(**processor_kwargs).input_ids,
593
594
595
596
597
598
599
                device=self.model.device.type,
            )

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

            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)
621
622
623
624
625
626
627
628
629
630
631

            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]

632
633
634
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

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

638
639
640
641
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
642
        del self.model
643
        cleanup_dist_env_and_memory()
644

Woosuk Kwon's avatar
Woosuk Kwon committed
645

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


class VllmRunner:

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

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

698
699
700
701
702
703
        if videos is not None:
            assert len(prompts) == len(videos)

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

704
705
706
        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
707
708
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
709

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

        if audios is not None:
            for i, audio in enumerate(audios):
Cyrus Leung's avatar
Cyrus Leung committed
717
718
                if audio is not None:
                    inputs[i]["multi_modal_data"] = {"audio": audio}
719
720
721

        return inputs

722
723
724
725
726
727
728
729
    def classify(self, prompts: List[str]) -> List[str]:
        req_outputs = self.model.encode(prompts)
        outputs = []
        for req_output in req_outputs:
            embedding = req_output.outputs.embedding
            outputs.append(embedding)
        return outputs

730
731
732
733
734
735
736
737
738
739
740
741
742
    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)

743
        req_outputs = self.model.generate(inputs,
744
                                          sampling_params=sampling_params)
745
746

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
747
748
749
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
750
751
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
752
753
            for sample in req_output.outputs:
                output_str = sample.text
754
                output_ids = list(sample.token_ids)
755
756
757
                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
758
759
        return outputs

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

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

789
        req_outputs = self.model.generate(inputs,
790
                                          sampling_params=sampling_params)
791
792
793
794
795
796
797

        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)
798
799
800

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

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

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

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

863
864
    def generate_encoder_decoder_greedy_logprobs(
        self,
865
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
866
867
        max_tokens: int,
        num_logprobs: int,
868
869
870
871
872
873
874
875
876
        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),
        )
877
878
879
880
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

881
        return self.generate_encoder_decoder_w_logprobs(
882
883
            encoder_decoder_prompts, greedy_logprobs_params)

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

Cyrus Leung's avatar
Cyrus Leung committed
900
901
902
903
904
905
906
907
908
909
910
911
912
913
    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)

        req_outputs = self.model.encode(inputs)
        return [req_output.outputs.embedding for req_output in req_outputs]
914

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

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

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

Woosuk Kwon's avatar
Woosuk Kwon committed
930

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


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={})
943
944
945
946
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
947
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963


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


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

971
    return cuda_device_count_stateless()
972
973
974


temp_dir = tempfile.gettempdir()
975
976
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
977
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "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
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
    return _dummy_opt_path


@pytest.fixture
def dummy_llava_path():
    json_path = os.path.join(_dummy_llava_path, "config.json")
    if not os.path.exists(_dummy_llava_path):
        snapshot_download(repo_id="llava-hf/llava-1.5-7b-hf",
                          local_dir=_dummy_llava_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1010
        with open(json_path) as f:
1011
1012
1013
1014
1015
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028


@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)
1029
        with open(json_path) as f:
1030
1031
1032
1033
1034
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053


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