conftest.py 36.2 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
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
youkaichao's avatar
youkaichao committed
31
32
                         TokensPrompt, to_enc_dec_tuple_list,
                         zip_enc_dec_prompts)
33
from vllm.logger import init_logger
34
from vllm.outputs import RequestOutput
35
from vllm.sampling_params import BeamSearchParams
36
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, cuda_device_count_stateless,
youkaichao's avatar
youkaichao committed
37
                        identity, is_list_of)
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
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
45

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

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

53

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


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


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

68
69

class _ImageAssets(_ImageAssetsBase):
70
71

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

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


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


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


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


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


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

    yield
Joe Runde's avatar
Joe Runde committed
127
128


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


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


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

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


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


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


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


182
183
184
185
186
187
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


188
189
190
191
192
193
194
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


195
@pytest.fixture
196
197
def example_encoder_decoder_prompts(
) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
198
199
200
201
202
203
    '''
    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:
204

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


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


236
237
238
239
240
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


241
242
243
244
245
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


246
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
247

Woosuk Kwon's avatar
Woosuk Kwon committed
248
249
250

class HfRunner:

251
    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
252
        from vllm.platforms import current_platform
253
254
255
        if x is None or isinstance(x, (bool, )):
            return x

256
        if device is None:
257
            device = "cpu" if current_platform.is_cpu() else "cuda"
258

259
260
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
261

262
263
264
265
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
266

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

281
        self.model_name = model_name
282

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

310
311
312
313
314
315
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
316

317
318
319
320
321
322
323
324
        # 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,
        )
325
326
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
327

328
        self.dtype = dtype
329
330
        self.postprocess_inputs = postprocess_inputs

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

341
342
343
344
345
346
347
        if videos is not None:
            assert len(prompts) == len(videos)

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

        all_inputs: List[BatchEncoding] = []
348
        for i, prompt in enumerate(prompts):
349
350
351
352
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
353
354
355
356
357
358
            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
359
360
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
361
362

            inputs = self.processor(**processor_kwargs)
363
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
364

365
366
367
368
            all_inputs.append(inputs)

        return all_inputs

369
370
371
372
373
374
375
376
377
378
379
    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

380
381
382
383
    def generate(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
384
        videos: Optional[PromptVideoInput] = None,
385
386
387
388
389
390
391
392
393
394
        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
395
            output_ids = self.model.generate(
396
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
397
398
399
                use_cache=True,
                **kwargs,
            )
400
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
401
402
403
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
404
405
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
406
407
408
409
410
411
412
            outputs.append((output_ids, output_str))
        return outputs

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

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
428
429
430
431
432
433

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
434
    ) -> List[Tuple[List[List[int]], List[str]]]:
435
436
437
438
439
440
441
442
443
444
445
446
447
448
        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
449

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

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

480
    def _hidden_states_to_seq_logprobs(
481
        self,
482
483
484
485
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

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

498
499
500
501
502
503
504
505
506
507
        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)

508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
        # 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,
        )

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

542
543
544
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
545

546
        for inputs in all_inputs:
547
            output = self.model.generate(
548
                **self.wrap_device(inputs, device=self.model.device.type),
549
550
551
552
553
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
554
                **kwargs,
555
556
            )

557
558
559
560
561
562
563
564
565
566
567
568
            (
                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))
569

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

586
587
588
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
589

590
591
592
593
594
595
596
597
        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]
598

599
            encoder_input_ids = self.wrap_device(
600
                self.processor(**processor_kwargs).input_ids,
601
602
603
604
605
606
607
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
608
                    self.tokenizer(decoder_prompt,
609
610
611
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628

            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)
629
630
631
632
633
634
635
636
637
638
639

            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]

640
641
642
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

643
644
645
    def predict(self, prompts: List[List[str]]) -> torch.Tensor:
        return self.model.predict(prompts, convert_to_tensor=True)

646
647
648
649
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
650
        del self.model
651
        cleanup_dist_env_and_memory()
652

Woosuk Kwon's avatar
Woosuk Kwon committed
653

Cyrus Leung's avatar
Cyrus Leung committed
654
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
655
656
657
658
659
660
661
662
663
def hf_runner():
    return HfRunner


class VllmRunner:

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

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

706
707
708
709
710
711
        if videos is not None:
            assert len(prompts) == len(videos)

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

712
713
714
        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
715
716
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
717

718
719
        if videos is not None:
            for i, video in enumerate(videos):
Cyrus Leung's avatar
Cyrus Leung committed
720
721
                if video is not None:
                    inputs[i]["multi_modal_data"] = {"video": video}
722
723
724

        if audios is not None:
            for i, audio in enumerate(audios):
Cyrus Leung's avatar
Cyrus Leung committed
725
726
                if audio is not None:
                    inputs[i]["multi_modal_data"] = {"audio": audio}
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742

        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)

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]]]:
youkaichao's avatar
youkaichao committed
890
891
892
893
894
895
        if is_list_of(prompts, str, check="all"):
            prompts = [TextPrompt(prompt=prompt) for prompt in prompts]
        else:
            prompts = [
                TokensPrompt(prompt_token_ids=tokens) for tokens in prompts
            ]
896
897
898
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
899
900
901
902
903
904
905
        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

906
907
908
909
    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
910
911
912
913
914
915
916
917
918
919
920
921
    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)

922
        req_outputs = self.model.embed(inputs)
Cyrus Leung's avatar
Cyrus Leung committed
923
        return [req_output.outputs.embedding for req_output in req_outputs]
924

925
926
927
928
    def score(
        self,
        text_1: Union[str, List[str]],
        text_2: Union[str, List[str]],
929
    ) -> List[float]:
930
        req_outputs = self.model.score(text_1, text_2)
931
        return [req_output.outputs.score for req_output in req_outputs]
932

933
934
935
936
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
937
        del self.model
938
        cleanup_dist_env_and_memory()
939

Woosuk Kwon's avatar
Woosuk Kwon committed
940

941
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
942
943
def vllm_runner():
    return VllmRunner
944
945
946
947
948
949
950
951
952


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={})
953
954
955
956
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
957
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973


@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
974
975
976
977
978
979
980


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

981
    return cuda_device_count_stateless()
982
983
984


temp_dir = tempfile.gettempdir()
985
986
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
987
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
988
989
990
991


@pytest.fixture
def dummy_opt_path():
992
993
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
994
        snapshot_download(repo_id="facebook/opt-125m",
995
                          local_dir=_dummy_opt_path,
996
997
998
999
1000
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1001
        with open(json_path) as f:
1002
1003
1004
1005
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
    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)
1020
        with open(json_path) as f:
1021
1022
1023
1024
1025
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038


@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)
1039
        with open(json_path) as f:
1040
1041
1042
1043
1044
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063


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