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

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

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

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

44
45
46
_TEST_DIR = os.path.dirname(__file__)
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
_LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
47

48
49
50
PromptImageInput = Union[List[Image.Image], List[List[Image.Image]]]
PromptAudioInput = Union[List[Tuple[np.ndarray, int]],
                         List[List[Tuple[np.ndarray, int]]]]
51
PromptVideoInput = Union[List[np.ndarray], List[List[np.ndarray]]]
52

53

54
def _read_prompts(filename: str) -> List[str]:
55
    with open(filename, "r") 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
67
68
69


if sys.version_info < (3, 9):
    # UserList cannot be subscripted
    class _ImageAssetsBase(UserList):
        pass
else:
70

71
72
    class _ImageAssetsBase(UserList[ImageAsset]):
        pass
73

74
75

class _ImageAssets(_ImageAssetsBase):
76
77

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

    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.
        """
90
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
91
92


93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


if sys.version_info < (3, 9):
    # UserList cannot be subscripted
    class _VideoAssetsBase(UserList):
        pass
else:

    class _VideoAssetsBase(UserList[VideoAsset]):
        pass


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


118
119
IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""
120
121
VIDEO_ASSETS = _VideoAssets()
"""Singleton instance of :class:`_VideoAssets`."""
122
123


124
125
126
127
128
129
130
@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


131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
@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
    cleanup()


146
147
def cleanup():
    destroy_model_parallel()
148
    destroy_distributed_environment()
149
150
151
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
152
153
    if not is_cpu():
        torch.cuda.empty_cache()
154
155


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

163
    return not request.node.get_closest_marker("skip_global_cleanup")
164
165


166
@pytest.fixture(autouse=True)
167
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
168
    yield
169
170
    if should_do_global_cleanup_after_test:
        cleanup()
171
172


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


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


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


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

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


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


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


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


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

Woosuk Kwon's avatar
Woosuk Kwon committed
247
248
249

class HfRunner:

250
251
252
253
254
255
256
257
    def wrap_device(self, input: _T, device: Optional[str] = None) -> _T:
        if device is None:
            return self.wrap_device(input, "cpu" if is_cpu() else "cuda")

        if hasattr(input, "device") and input.device.type == device:
            return input

        return input.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
        auto_cls: Type[_BaseAutoModelClass] = AutoModelForCausalLM,
267
268
        postprocess_inputs: Callable[[BatchEncoding],
                                     BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
269
    ) -> None:
270
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
271

272
        self.model_name = model_name
273

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

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
        )

299
300
301
302
303
304
305
306
        # 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,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
307

308
309
        self.postprocess_inputs = postprocess_inputs

310
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
311
312
        self,
        prompts: List[str],
313
        images: Optional[PromptImageInput] = None,
314
315
316
317
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[BatchEncoding]:
        if images is not None:
318
            assert len(prompts) == len(images)
319

320
321
322
323
324
325
326
        if videos is not None:
            assert len(prompts) == len(videos)

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

        all_inputs: List[BatchEncoding] = []
327
        for i, prompt in enumerate(prompts):
328
329
330
331
332
333
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
334
335
            if videos is not None and videos[i] is not None:
                processor_kwargs["videos"] = videos[i]
336
337
338
339
            if audios is not None and audios[i] is not None:
                audio, sr = audios[i]
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
340
341

            inputs = self.processor(**processor_kwargs)
342
            inputs = self.postprocess_inputs(inputs)
343

344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
            all_inputs.append(inputs)

        return all_inputs

    def generate(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
        videos: Optional[List[np.ndarray]] = None,
        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
363
            output_ids = self.model.generate(
364
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
365
366
367
                use_cache=True,
                **kwargs,
            )
368
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
369
370
371
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
372
373
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
374
375
376
377
378
379
380
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
381
        images: Optional[PromptImageInput] = None,
382
383
        videos: Optional[List[np.ndarray]] = None,
        audios: Optional[PromptAudioInput] = None,
384
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
385
    ) -> List[Tuple[List[int], str]]:
386
387
        outputs = self.generate(prompts,
                                do_sample=False,
388
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
389
                                images=images,
390
391
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
392
                                **kwargs)
393
394
395

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
396
397
398
399
400
401

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
402
    ) -> List[Tuple[List[List[int]], List[str]]]:
403
404
405
406
407
408
409
410
411
412
413
414
415
416
        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
417

418
419
420
421
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
422
        images: Optional[PromptImageInput] = None,
423
        videos: Optional[List[np.ndarray]] = None,
424
        audios: Optional[PromptAudioInput] = None,
425
        **kwargs: Any,
426
    ) -> List[List[torch.Tensor]]:
427
428
429
430
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
431

432
433
        all_logprobs: List[List[torch.Tensor]] = []
        for inputs in all_inputs:
434
            output = self.model.generate(
435
                **self.wrap_device(inputs, device=self.model.device.type),
436
437
438
439
440
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
441
                **kwargs,
442
            )
443
444
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
445
446
447
            all_logprobs.append(seq_logprobs)
        return all_logprobs

448
    def _hidden_states_to_seq_logprobs(
449
        self,
450
451
452
453
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

454
455
456
457
        seq_logprobs: List[torch.Tensor] = []
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
458
459
                last_hidden_states.to(output_embeddings.weight.device),
                output_embeddings.weight.t(),
460
            )
461
462
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
463
464
465
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

466
467
468
469
470
471
472
473
474
475
        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)

476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
        # 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,
        )

495
496
497
498
499
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
500
501
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
502
        videos: Optional[List[np.ndarray]] = None,
503
        **kwargs: Any,
504
    ) -> List[TokensTextLogprobs]:
505
506
507
508
509
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

510
511
512
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
513

514
        for inputs in all_inputs:
515
            output = self.model.generate(
516
                **self.wrap_device(inputs, device=self.model.device.type),
517
518
519
520
521
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
522
                **kwargs,
523
524
            )

525
526
527
528
529
530
531
532
533
534
535
536
            (
                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))
537

538
539
540
541
542
543
        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,
544
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
545
546
547
        max_tokens: int,
        num_logprobs: int,
        **kwargs: Any,
548
    ) -> List[TokensTextLogprobs]:
549
550
551
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
552

553
554
555
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
556

557
558
        for (encoder_prompt,
             decoder_prompt) in to_enc_dec_tuple_list(encoder_decoder_prompts):
559

560
            encoder_input_ids = self.wrap_device(
561
562
563
564
565
566
567
568
                self.tokenizer(encoder_prompt, return_tensors="pt").input_ids,
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
569
                    self.tokenizer(decoder_prompt,
570
571
572
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589

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

            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]

601
602
603
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

604
605
606
607
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
608
609
610
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
611

Cyrus Leung's avatar
Cyrus Leung committed
612
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
613
614
615
616
617
618
619
620
621
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
622
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
623
        tokenizer_name: Optional[str] = None,
624
625
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
626
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
627
        dtype: str = "half",
628
        disable_log_stats: bool = True,
629
        tensor_parallel_size: int = 1,
630
631
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
632
        swap_space: int = 4,
633
        enforce_eager: Optional[bool] = False,
634
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
635
636
637
    ) -> None:
        self.model = LLM(
            model=model_name,
638
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
639
640
641
            tokenizer=tokenizer_name,
            trust_remote_code=True,
            dtype=dtype,
642
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
643
            enforce_eager=enforce_eager,
644
            disable_log_stats=disable_log_stats,
645
            tensor_parallel_size=tensor_parallel_size,
646
            max_model_len=max_model_len,
647
648
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
649
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
650
651
        )

652
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
653
654
        self,
        prompts: List[str],
655
        images: Optional[PromptImageInput] = None,
656
657
658
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[TextPrompt]:
659
        if images is not None:
660
            assert len(prompts) == len(images)
661

662
663
664
665
666
667
        if videos is not None:
            assert len(prompts) == len(videos)

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

668
669
670
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
671
                inputs[i]["multi_modal_data"] = {"image": image}
672

673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
        if videos is not None:
            for i, video in enumerate(videos):
                inputs[i]["multi_modal_data"] = {"video": video}

        if audios is not None:
            for i, audio in enumerate(audios):
                inputs[i]["multi_modal_data"] = {"audio": audio}

        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)

696
        req_outputs = self.model.generate(inputs,
697
                                          sampling_params=sampling_params)
698
699

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
700
701
702
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
703
704
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
705
706
            for sample in req_output.outputs:
                output_str = sample.text
707
                output_ids = list(sample.token_ids)
708
709
710
                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
711
712
        return outputs

713
    @staticmethod
714
715
    def _final_steps_generate_w_logprobs(
        req_outputs: List[RequestOutput],
716
717
    ) -> List[TokensTextLogprobsPromptLogprobs]:
        outputs: List[TokensTextLogprobsPromptLogprobs] = []
718
        for req_output in req_outputs:
719
            assert len(req_output.outputs) > 0
720
721
            for sample in req_output.outputs:
                output_str = sample.text
722
                output_ids = list(sample.token_ids)
723
                output_logprobs = sample.logprobs
724
725
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
726
727
        return outputs

728
729
730
731
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
732
733
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
734
        videos: Optional[PromptVideoInput] = None,
735
736
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
737
738
739
740
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
741

742
        req_outputs = self.model.generate(inputs,
743
                                          sampling_params=sampling_params)
744
745
746
747
748
749
750

        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)
751
752
753

    def generate_encoder_decoder_w_logprobs(
        self,
754
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
755
        sampling_params: SamplingParams,
756
757
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
758
759
760
761
762
763
764
        '''
        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)
765
766
767
768
769
770
        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)
771

Woosuk Kwon's avatar
Woosuk Kwon committed
772
773
774
775
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
776
        images: Optional[PromptImageInput] = None,
777
778
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
779
780
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
781
782
783
784
785
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
                                audios=audios)
786
787
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
788

789
790
791
792
793
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
794
        num_prompt_logprobs: Optional[int] = None,
795
796
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
797
        videos: Optional[PromptVideoInput] = None,
798
        stop_token_ids: Optional[List[int]] = None,
799
800
801
802
803
804
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
805
            prompt_logprobs=num_prompt_logprobs,
806
807
808
809
810
811
812
            stop_token_ids=stop_token_ids)

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

814
815
    def generate_encoder_decoder_greedy_logprobs(
        self,
816
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
817
818
        max_tokens: int,
        num_logprobs: int,
819
820
821
822
823
824
825
826
827
        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),
        )
828
829
830
831
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

832
        return self.generate_encoder_decoder_w_logprobs(
833
834
            encoder_decoder_prompts, greedy_logprobs_params)

835
    def generate_beam_search(
836
837
838
839
840
        self,
        prompts: Union[List[str], List[List[int]]],
        beam_width: int,
        max_tokens: int,
    ) -> List[Tuple[List[List[int]], List[str]]]:
841
842
843
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
844
845
846
847
848
849
850
        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

851
852
853
854
855
856
857
858
    def encode(self, prompts: List[str]) -> List[List[float]]:
        req_outputs = self.model.encode(prompts)
        outputs = []
        for req_output in req_outputs:
            embedding = req_output.outputs.embedding
            outputs.append(embedding)
        return outputs

859
860
861
862
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
863
864
865
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
866

867
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
868
869
def vllm_runner():
    return VllmRunner
870
871
872
873
874
875
876
877
878


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={})
879
880
881
882
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
883
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899


@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
900
901
902
903
904
905
906


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

907
    return cuda_device_count_stateless()
908
909
910


temp_dir = tempfile.gettempdir()
911
912
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
913
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
914
915
916
917


@pytest.fixture
def dummy_opt_path():
918
919
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
920
        snapshot_download(repo_id="facebook/opt-125m",
921
                          local_dir=_dummy_opt_path,
922
923
924
925
926
927
928
929
930
931
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
        with open(json_path, "r") as f:
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
    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)
        with open(json_path, "r") as f:
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970


@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)
        with open(json_path, "r") as f:
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path