conftest.py 36.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import json
4
import os
5
import tempfile
6
from collections import UserList
7
from enum import Enum
8
9
from typing import (Any, Callable, Dict, List, Optional, Tuple, Type,
                    TypedDict, TypeVar, Union)
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
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
youkaichao's avatar
youkaichao committed
33
34
                         TokensPrompt, to_enc_dec_tuple_list,
                         zip_enc_dec_prompts)
35
from vllm.logger import init_logger
36
from vllm.outputs import RequestOutput
37
from vllm.sampling_params import BeamSearchParams
38
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, cuda_device_count_stateless,
youkaichao's avatar
youkaichao committed
39
                        identity, is_list_of)
40

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

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

Cyrus Leung's avatar
Cyrus Leung committed
48
_M = TypeVar("_M")
49

Cyrus Leung's avatar
Cyrus Leung committed
50
51
52
53
54
_PromptMultiModalInput = Union[List[_M], List[List[_M]]]

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

56

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


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


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

71
72

class _ImageAssets(_ImageAssetsBase):
73
74

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

    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.
        """
87
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
88
89


90
91
92
93
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


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


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


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


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

    yield
Joe Runde's avatar
Joe Runde committed
130
131


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


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


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

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


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


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


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


185
186
187
188
189
190
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


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


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

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


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


239
240
241
242
243
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


244
245
246
247
248
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


249
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
250
_R = TypeVar("_R")
251

Woosuk Kwon's avatar
Woosuk Kwon committed
252
253
254

class HfRunner:

255
    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
256
        from vllm.platforms import current_platform
257
258
259
        if x is None or isinstance(x, (bool, )):
            return x

260
        if device is None:
261
            device = "cpu" if current_platform.is_cpu() else "cuda"
262

263
264
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
265

266
267
268
269
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
270

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

285
        self.model_name = model_name
286

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

314
315
316
317
318
319
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
320

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

332
        self.dtype = dtype
333
334
        self.postprocess_inputs = postprocess_inputs

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

345
346
347
348
349
350
351
        if videos is not None:
            assert len(prompts) == len(videos)

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

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

            inputs = self.processor(**processor_kwargs)
367
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
368

369
370
371
372
            all_inputs.append(inputs)

        return all_inputs

373
374
375
376
377
378
379
380
381
382
383
    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

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

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

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
432
433
434
435
436
437

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

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

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

484
    def _hidden_states_to_seq_logprobs(
485
        self,
486
487
488
489
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

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

502
503
504
505
506
507
508
509
510
511
        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)

512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        # 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,
        )

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

546
547
548
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
549

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

561
562
563
564
565
566
567
568
569
570
571
572
            (
                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))
573

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

590
591
592
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
593

594
595
596
597
598
599
600
601
        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]
602

603
604
            encoder_inputs = self.wrap_device(
                self.processor(**processor_kwargs),
605
606
607
608
609
610
611
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
612
                    self.tokenizer(decoder_prompt,
613
614
615
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
616
617
618
619
620
621
622
623

            output = self.model.generate(
                decoder_input_ids=decoder_input_ids,
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
624
                **encoder_inputs,
625
626
627
628
629
630
631
632
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
633
634
635
636
637
638
639
640
641
642
643

            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]

644
645
646
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

647
648
649
    def predict(self, prompts: List[List[str]]) -> torch.Tensor:
        return self.model.predict(prompts, convert_to_tensor=True)

650
651
652
653
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
654
        del self.model
655
        cleanup_dist_env_and_memory()
656

Woosuk Kwon's avatar
Woosuk Kwon committed
657

Cyrus Leung's avatar
Cyrus Leung committed
658
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
659
660
661
662
663
664
665
666
667
def hf_runner():
    return HfRunner


class VllmRunner:

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

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

710
711
712
713
714
715
        if videos is not None:
            assert len(prompts) == len(videos)

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

716
717
718
        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
719
720
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
721

722
723
        if videos is not None:
            for i, video in enumerate(videos):
Cyrus Leung's avatar
Cyrus Leung committed
724
725
                if video is not None:
                    inputs[i]["multi_modal_data"] = {"video": video}
726
727
728

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

        return inputs

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
741
        **kwargs: Any,
742
743
744
745
746
747
    ) -> List[Tuple[List[List[int]], List[str]]]:
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

748
        req_outputs = self.model.generate(inputs,
749
750
                                          sampling_params=sampling_params,
                                          **kwargs)
751
752

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

766
    @staticmethod
767
768
    def _final_steps_generate_w_logprobs(
        req_outputs: List[RequestOutput],
769
770
    ) -> List[TokensTextLogprobsPromptLogprobs]:
        outputs: List[TokensTextLogprobsPromptLogprobs] = []
771
        for req_output in req_outputs:
772
            assert len(req_output.outputs) > 0
773
774
            for sample in req_output.outputs:
                output_str = sample.text
775
                output_ids = list(sample.token_ids)
776
                output_logprobs = sample.logprobs
777
778
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
779
780
        return outputs

781
782
783
784
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
785
786
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
787
        videos: Optional[PromptVideoInput] = None,
788
        **kwargs: Any,
789
790
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
791
792
793
794
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
795

796
        req_outputs = self.model.generate(inputs,
797
798
                                          sampling_params=sampling_params,
                                          **kwargs)
799
800
801
802
803
804
805

        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)
806
807
808

    def generate_encoder_decoder_w_logprobs(
        self,
809
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
810
        sampling_params: SamplingParams,
811
812
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
813
814
815
816
817
818
819
        '''
        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)
820
821
822
823
824
825
        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)
826

Woosuk Kwon's avatar
Woosuk Kwon committed
827
828
829
830
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
831
        images: Optional[PromptImageInput] = None,
832
833
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
834
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
835
836
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
837
838
839
840
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
841
842
                                audios=audios,
                                **kwargs)
843
844
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
845

846
847
848
849
850
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
851
        num_prompt_logprobs: Optional[int] = None,
852
853
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
854
        videos: Optional[PromptVideoInput] = None,
855
        stop_token_ids: Optional[List[int]] = None,
856
        stop: Optional[List[str]] = None,
857
        **kwargs: Any,
858
859
860
861
862
863
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
864
            prompt_logprobs=num_prompt_logprobs,
865
866
            stop_token_ids=stop_token_ids,
            stop=stop)
867
868
869
870
871

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
872
873
                                        videos=videos,
                                        **kwargs)
874

875
876
    def generate_encoder_decoder_greedy_logprobs(
        self,
877
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
878
879
        max_tokens: int,
        num_logprobs: int,
880
881
882
883
884
885
886
887
888
        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),
        )
889
890
891
892
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

893
        return self.generate_encoder_decoder_w_logprobs(
894
895
            encoder_decoder_prompts, greedy_logprobs_params)

896
    def generate_beam_search(
897
898
899
900
901
        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
902
903
904
905
906
907
        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
            ]
908
909
910
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
911
912
913
914
915
916
917
        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

918
919
920
921
    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
922
923
924
925
926
927
928
929
930
931
932
933
    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)

934
        req_outputs = self.model.embed(inputs)
Cyrus Leung's avatar
Cyrus Leung committed
935
        return [req_output.outputs.embedding for req_output in req_outputs]
936

937
938
939
940
    def score(
        self,
        text_1: Union[str, List[str]],
        text_2: Union[str, List[str]],
941
    ) -> List[float]:
942
        req_outputs = self.model.score(text_1, text_2)
943
        return [req_output.outputs.score for req_output in req_outputs]
944

945
946
947
948
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
        executor = self.model.llm_engine.model_executor
        return executor.apply_model(func)

949
950
951
952
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
953
        del self.model
954
        cleanup_dist_env_and_memory()
955

Woosuk Kwon's avatar
Woosuk Kwon committed
956

957
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
958
959
def vllm_runner():
    return VllmRunner
960
961
962
963
964
965
966
967
968


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={})
969
970
971
972
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
973
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989


@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
990
991
992
993
994
995
996


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

997
    return cuda_device_count_stateless()
998
999
1000


temp_dir = tempfile.gettempdir()
1001
1002
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1003
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1004
1005
1006
1007


@pytest.fixture
def dummy_opt_path():
1008
1009
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1010
        snapshot_download(repo_id="facebook/opt-125m",
1011
                          local_dir=_dummy_opt_path,
1012
1013
1014
1015
1016
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1017
        with open(json_path) as f:
1018
1019
1020
1021
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
    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)
1036
        with open(json_path) as f:
1037
1038
1039
1040
1041
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054


@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)
1055
        with open(json_path) as f:
1056
1057
1058
1059
1060
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079


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