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

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

21
22
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
23
from vllm import LLM, SamplingParams
24
from vllm.assets.image import ImageAsset
25
from vllm.assets.video import VideoAsset
26
from vllm.config import TaskOption, TokenizerPoolConfig
27
from vllm.connections import global_http_connection
28
from vllm.distributed import (cleanup_dist_env_and_memory,
29
30
                              init_distributed_environment,
                              initialize_model_parallel)
31
32
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
                         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.platforms import current_platform
36
from vllm.sampling_params import BeamSearchParams
37
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, cuda_device_count_stateless,
38
                        identity)
39

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

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

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

51

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


58
59
60
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
61
62
63
64
65
66
67


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

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

72
73

class _ImageAssets(_ImageAssetsBase):
74
75

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

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


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


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


122
123
124
125
126
127
128
@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


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


144
@pytest.fixture()
145
def should_do_global_cleanup_after_test(request) -> bool:
146
147
148
149
    """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.
    """
150

151
    return not request.node.get_closest_marker("skip_global_cleanup")
152
153


154
@pytest.fixture(autouse=True)
155
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
156
    yield
157
    if should_do_global_cleanup_after_test:
158
        cleanup_dist_env_and_memory()
159
160


161
162
163
164
165
166
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
167
168
@pytest.fixture
def example_prompts() -> List[str]:
169
170
    prompts = []
    for filename in _TEST_PROMPTS:
171
        prompts += _read_prompts(filename)
172
173
174
    return prompts


175
176
177
178
179
180
181
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


182
@pytest.fixture
183
184
def example_encoder_decoder_prompts(
) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
185
186
187
188
189
190
    '''
    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:
191

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
    * 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:
207
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
208
        DecoderPromptType.EMPTY_STR:
209
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
210
        DecoderPromptType.CUSTOM:
211
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
212
213
214
    }


215
216
217
218
@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
219
        prompts += _read_prompts(filename)
220
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
221
222


223
224
225
226
227
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


228
229
230
231
232
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


233
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature)
234

Woosuk Kwon's avatar
Woosuk Kwon committed
235
236
237

class HfRunner:

238
239
    def wrap_device(self, input: _T, device: Optional[str] = None) -> _T:
        if device is None:
240
241
            return self.wrap_device(
                input, "cpu" if current_platform.is_cpu() else "cuda")
242
243
244
245
246

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

        return input.to(device)
247

Woosuk Kwon's avatar
Woosuk Kwon committed
248
249
250
251
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
252
        *,
253
        model_kwargs: Optional[Dict[str, Any]] = None,
254
        is_sentence_transformer: bool = False,
255
        auto_cls: Type[_BaseAutoModelClass] = AutoModelForCausalLM,
256
257
        postprocess_inputs: Callable[[BatchEncoding],
                                     BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
258
    ) -> None:
259
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
260

261
        self.model_name = model_name
262

263
        if is_sentence_transformer:
264
265
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
266
267
268
269
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
270
                    trust_remote_code=True,
271
                ).to(dtype=torch_dtype))
272
        else:
273
            model_kwargs = model_kwargs if model_kwargs is not None else {}
274
            self.model = self.wrap_device(
275
                auto_cls.from_pretrained(
276
277
278
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
279
                    **model_kwargs,
280
                ))
281
282
283
284
285
286
287

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

288
289
290
291
292
293
294
295
        # 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
296

297
298
        self.postprocess_inputs = postprocess_inputs

299
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
300
301
        self,
        prompts: List[str],
302
        images: Optional[PromptImageInput] = None,
303
304
305
306
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[BatchEncoding]:
        if images is not None:
307
            assert len(prompts) == len(images)
308

309
310
311
312
313
314
315
        if videos is not None:
            assert len(prompts) == len(videos)

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

        all_inputs: List[BatchEncoding] = []
316
        for i, prompt in enumerate(prompts):
317
318
319
320
321
322
            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]
323
324
            if videos is not None and videos[i] is not None:
                processor_kwargs["videos"] = videos[i]
325
326
327
328
            if audios is not None and audios[i] is not None:
                audio, sr = audios[i]
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
329
330

            inputs = self.processor(**processor_kwargs)
331
            inputs = self.postprocess_inputs(inputs)
332

333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
            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
352
            output_ids = self.model.generate(
353
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
354
355
356
                use_cache=True,
                **kwargs,
            )
357
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
358
359
360
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
361
362
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
363
364
365
366
367
368
369
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
370
        images: Optional[PromptImageInput] = None,
371
372
        videos: Optional[List[np.ndarray]] = None,
        audios: Optional[PromptAudioInput] = None,
373
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
374
    ) -> List[Tuple[List[int], str]]:
375
376
        outputs = self.generate(prompts,
                                do_sample=False,
377
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
378
                                images=images,
379
380
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
381
                                **kwargs)
382
383
384

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
385
386
387
388
389
390

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
391
    ) -> List[Tuple[List[List[int]], List[str]]]:
392
393
394
395
396
397
398
399
400
401
402
403
404
405
        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
406

407
408
409
410
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
411
        images: Optional[PromptImageInput] = None,
412
        videos: Optional[List[np.ndarray]] = None,
413
        audios: Optional[PromptAudioInput] = None,
414
        **kwargs: Any,
415
    ) -> List[List[torch.Tensor]]:
416
417
418
419
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
420

421
422
        all_logprobs: List[List[torch.Tensor]] = []
        for inputs in all_inputs:
423
            output = self.model.generate(
424
                **self.wrap_device(inputs, device=self.model.device.type),
425
426
427
428
429
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
430
                **kwargs,
431
            )
432
433
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
434
435
436
            all_logprobs.append(seq_logprobs)
        return all_logprobs

437
    def _hidden_states_to_seq_logprobs(
438
        self,
439
440
441
442
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

443
444
445
446
        seq_logprobs: List[torch.Tensor] = []
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
447
448
                last_hidden_states.to(output_embeddings.weight.device),
                output_embeddings.weight.t(),
449
            )
450
451
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
452
453
454
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

455
456
457
458
459
460
461
462
463
464
        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)

465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
        # 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,
        )

484
485
486
487
488
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
489
490
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
491
        videos: Optional[List[np.ndarray]] = None,
492
        **kwargs: Any,
493
    ) -> List[TokensTextLogprobs]:
494
495
496
497
498
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

499
500
501
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
502

503
        for inputs in all_inputs:
504
            output = self.model.generate(
505
                **self.wrap_device(inputs, device=self.model.device.type),
506
507
508
509
510
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
511
                **kwargs,
512
513
            )

514
515
516
517
518
519
520
521
522
523
524
525
            (
                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))
526

527
528
529
530
531
532
        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,
533
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
534
535
536
        max_tokens: int,
        num_logprobs: int,
        **kwargs: Any,
537
    ) -> List[TokensTextLogprobs]:
538
539
540
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
541

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

546
547
        for (encoder_prompt,
             decoder_prompt) in to_enc_dec_tuple_list(encoder_decoder_prompts):
548

549
            encoder_input_ids = self.wrap_device(
550
551
552
553
554
555
556
557
                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(
558
                    self.tokenizer(decoder_prompt,
559
560
561
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578

            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)
579
580
581
582
583
584
585
586
587
588
589

            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]

590
591
592
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

593
594
595
596
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
597
        del self.model
598
        cleanup_dist_env_and_memory()
599

Woosuk Kwon's avatar
Woosuk Kwon committed
600

Cyrus Leung's avatar
Cyrus Leung committed
601
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
602
603
604
605
606
607
608
609
610
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
611
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
612
        tokenizer_name: Optional[str] = None,
613
614
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
615
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
616
        dtype: str = "half",
617
        disable_log_stats: bool = True,
618
        tensor_parallel_size: int = 1,
619
620
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
621
        swap_space: int = 4,
622
        enforce_eager: Optional[bool] = False,
623
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
624
625
626
    ) -> None:
        self.model = LLM(
            model=model_name,
627
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
628
629
630
            tokenizer=tokenizer_name,
            trust_remote_code=True,
            dtype=dtype,
631
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
632
            enforce_eager=enforce_eager,
633
            disable_log_stats=disable_log_stats,
634
            tensor_parallel_size=tensor_parallel_size,
635
            max_model_len=max_model_len,
636
637
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
638
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
639
640
        )

641
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
642
643
        self,
        prompts: List[str],
644
        images: Optional[PromptImageInput] = None,
645
646
647
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[TextPrompt]:
648
        if images is not None:
649
            assert len(prompts) == len(images)
650

651
652
653
654
655
656
        if videos is not None:
            assert len(prompts) == len(videos)

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

657
658
659
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
660
                inputs[i]["multi_modal_data"] = {"image": image}
661

662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
        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)

685
        req_outputs = self.model.generate(inputs,
686
                                          sampling_params=sampling_params)
687
688

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
689
690
691
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
692
693
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
694
695
            for sample in req_output.outputs:
                output_str = sample.text
696
                output_ids = list(sample.token_ids)
697
698
699
                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
700
701
        return outputs

702
    @staticmethod
703
704
    def _final_steps_generate_w_logprobs(
        req_outputs: List[RequestOutput],
705
706
    ) -> List[TokensTextLogprobsPromptLogprobs]:
        outputs: List[TokensTextLogprobsPromptLogprobs] = []
707
        for req_output in req_outputs:
708
            assert len(req_output.outputs) > 0
709
710
            for sample in req_output.outputs:
                output_str = sample.text
711
                output_ids = list(sample.token_ids)
712
                output_logprobs = sample.logprobs
713
714
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
715
716
        return outputs

717
718
719
720
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
721
722
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
723
        videos: Optional[PromptVideoInput] = None,
724
725
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
726
727
728
729
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
730

731
        req_outputs = self.model.generate(inputs,
732
                                          sampling_params=sampling_params)
733
734
735
736
737
738
739

        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)
740
741
742

    def generate_encoder_decoder_w_logprobs(
        self,
743
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
744
        sampling_params: SamplingParams,
745
746
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
747
748
749
750
751
752
753
        '''
        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)
754
755
756
757
758
759
        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)
760

Woosuk Kwon's avatar
Woosuk Kwon committed
761
762
763
764
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
765
        images: Optional[PromptImageInput] = None,
766
767
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
768
769
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
770
771
772
773
774
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
                                audios=audios)
775
776
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
777

778
779
780
781
782
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
783
        num_prompt_logprobs: Optional[int] = None,
784
785
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
786
        videos: Optional[PromptVideoInput] = None,
787
        stop_token_ids: Optional[List[int]] = None,
788
789
790
791
792
793
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
794
            prompt_logprobs=num_prompt_logprobs,
795
796
797
798
799
800
801
            stop_token_ids=stop_token_ids)

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

803
804
    def generate_encoder_decoder_greedy_logprobs(
        self,
805
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
806
807
        max_tokens: int,
        num_logprobs: int,
808
809
810
811
812
813
814
815
816
        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),
        )
817
818
819
820
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

821
        return self.generate_encoder_decoder_w_logprobs(
822
823
            encoder_decoder_prompts, greedy_logprobs_params)

824
    def generate_beam_search(
825
826
827
828
829
        self,
        prompts: Union[List[str], List[List[int]]],
        beam_width: int,
        max_tokens: int,
    ) -> List[Tuple[List[List[int]], List[str]]]:
830
831
832
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
833
834
835
836
837
838
839
        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

840
841
842
843
844
845
846
847
    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

848
849
850
851
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
852
        del self.model
853
        cleanup_dist_env_and_memory()
854

Woosuk Kwon's avatar
Woosuk Kwon committed
855

856
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
857
858
def vllm_runner():
    return VllmRunner
859
860
861
862
863
864
865
866
867


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={})
868
869
870
871
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
872
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888


@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
889
890
891
892
893
894
895


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

896
    return cuda_device_count_stateless()
897
898
899


temp_dir = tempfile.gettempdir()
900
901
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
902
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
903
904
905
906


@pytest.fixture
def dummy_opt_path():
907
908
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
909
        snapshot_download(repo_id="facebook/opt-125m",
910
                          local_dir=_dummy_opt_path,
911
912
913
914
915
916
917
918
919
920
                          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)
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
    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
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959


@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