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

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

49
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
50
51

PromptImageInput = _PromptMultiModalInput[Image.Image]
52
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
53
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
54

55

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


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


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

70
71

class _ImageAssets(_ImageAssetsBase):
72
73

    def __init__(self) -> None:
74
75
76
77
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
78

79
    def prompts(self, prompts: _ImageAssetPrompts) -> list[str]:
80
81
82
83
84
85
        """
        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.
        """
86
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
87
88


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


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


class _VideoAssets(_VideoAssetsBase):

    def __init__(self) -> None:
        super().__init__([
            VideoAsset("sample_demo_1.mp4"),
        ])

104
    def prompts(self, prompts: _VideoAssetPrompts) -> list[str]:
105
106
107
        return [prompts["sample_demo_1"]]


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


114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
@pytest.fixture(scope="function", autouse=True)
def cleanup_VLLM_USE_V1(monkeypatch):
    """
    The V1 oracle sets "VLLM_USE_V1" during loading. This means
    that each invocation of a test change the env variable.

    If we touch "VLLM_USE_V1" with monkeypatch, then any changes
    made during the test run by vLLM will be cleaned up.

    This fixture is used by every test.
    """

    # If VLLM_USE_V1 is not set, set then delete. This will
    # cause monkeypatch to clean up VLLM_USE_V1 upon exit
    # if VLLM modifies the value of envs.VLLM_USE_V1.
    if "VLLM_USE_V1" not in os.environ:
        monkeypatch.setenv("VLLM_USE_V1", "")
        monkeypatch.delenv("VLLM_USE_V1")


Joe Runde's avatar
Joe Runde committed
134
@pytest.fixture(params=[True, False])
135
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
136
137
138
139
140
141
142
143
    # 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")
144
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
145
    else:
146
147
148
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
149
150


151
152
153
154
155
156
157
@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


158
159
160
161
162
163
164
165
166
167
168
169
@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
170
    cleanup_dist_env_and_memory()
171
172


173
@pytest.fixture()
174
def should_do_global_cleanup_after_test(request) -> bool:
175
176
177
178
    """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.
    """
179

180
    return not request.node.get_closest_marker("skip_global_cleanup")
181
182


183
@pytest.fixture(autouse=True)
184
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
185
    yield
186
    if should_do_global_cleanup_after_test:
187
        cleanup_dist_env_and_memory()
188
189


190
191
192
193
194
195
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
196
@pytest.fixture
197
def example_prompts() -> list[str]:
198
199
    prompts = []
    for filename in _TEST_PROMPTS:
200
        prompts += _read_prompts(filename)
201
202
203
    return prompts


204
205
206
207
208
209
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


210
211
212
213
214
215
216
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


217
@pytest.fixture
218
def example_encoder_decoder_prompts(
219
) -> dict[DecoderPromptType, list[ExplicitEncoderDecoderPrompt]]:
220
221
222
223
224
225
    '''
    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:
226

227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    * 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:
242
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
243
        DecoderPromptType.EMPTY_STR:
244
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
245
        DecoderPromptType.CUSTOM:
246
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
247
248
249
    }


250
@pytest.fixture
251
def example_long_prompts() -> list[str]:
252
253
    prompts = []
    for filename in _LONG_PROMPTS:
254
        prompts += _read_prompts(filename)
255
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
256
257


258
259
260
261
262
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


263
264
265
266
267
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


268
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
269
_R = TypeVar("_R")
270

Woosuk Kwon's avatar
Woosuk Kwon committed
271
272
273

class HfRunner:

274
    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
275
        from vllm.platforms import current_platform
276
277
278
        if x is None or isinstance(x, (bool, )):
            return x

279
        if device is None:
280
281
            device = "cpu" if current_platform.is_cpu(
            ) or current_platform.is_openvino() else "cuda"
282

283
284
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
285

286
287
288
289
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
290

Woosuk Kwon's avatar
Woosuk Kwon committed
291
292
293
294
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
295
        *,
296
        model_kwargs: Optional[dict[str, Any]] = None,
297
        is_sentence_transformer: bool = False,
298
        is_cross_encoder: bool = False,
299
        skip_tokenizer_init: bool = False,
300
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
301
        postprocess_inputs: Callable[..., BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
302
    ) -> None:
303
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
304

305
        self.model_name = model_name
306

307
        if is_sentence_transformer:
308
309
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
310
311
312
313
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
314
                    trust_remote_code=True,
315
                ).to(dtype=torch_dtype))
316
317
318
319
320
321
322
323
        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)
324
        else:
325
            model_kwargs = model_kwargs if model_kwargs is not None else {}
326
            self.model = self.wrap_device(
327
                auto_cls.from_pretrained(
328
329
330
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
331
                    **model_kwargs,
332
                ))
333

334
335
336
337
338
339
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
340

341
342
343
344
345
346
347
348
        # 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,
        )
349
350
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
351

352
        self.dtype = dtype
353
354
        self.postprocess_inputs = postprocess_inputs

355
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
356
        self,
357
        prompts: list[str],
358
        images: Optional[PromptImageInput] = None,
359
360
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
361
    ) -> list[BatchEncoding]:
362
        if images is not None:
363
            assert len(prompts) == len(images)
364

365
366
367
368
369
370
        if videos is not None:
            assert len(prompts) == len(videos)

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

371
        all_inputs: list[BatchEncoding] = []
372
        for i, prompt in enumerate(prompts):
373
            processor_kwargs: dict[str, Any] = {
374
375
376
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
377
378
379
380
381
382
            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
383
384
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
385
386

            inputs = self.processor(**processor_kwargs)
387
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
388

389
390
391
392
            all_inputs.append(inputs)

        return all_inputs

393
    def classify(self, prompts: list[str]) -> list[str]:
394
395
396
397
398
399
400
401
402
403
        # 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

404
405
    def generate(
        self,
406
        prompts: list[str],
407
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
408
        videos: Optional[PromptVideoInput] = None,
409
410
        audios: Optional[PromptAudioInput] = None,
        **kwargs: Any,
411
    ) -> list[tuple[list[list[int]], list[str]]]:
412
413
414
415
416
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

417
        outputs: list[tuple[list[list[int]], list[str]]] = []
418
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
419
            output_ids = self.model.generate(
420
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
421
422
423
                use_cache=True,
                **kwargs,
            )
424
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
425
426
427
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
428
429
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
430
431
432
433
434
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
435
        prompts: list[str],
Woosuk Kwon's avatar
Woosuk Kwon committed
436
        max_tokens: int,
437
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
438
        videos: Optional[PromptVideoInput] = None,
439
        audios: Optional[PromptAudioInput] = None,
440
        **kwargs: Any,
441
    ) -> list[tuple[list[int], str]]:
442
443
        outputs = self.generate(prompts,
                                do_sample=False,
444
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
445
                                images=images,
446
447
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
448
                                **kwargs)
449
450
451

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
452
453
454

    def generate_beam_search(
        self,
455
        prompts: list[str],
456
457
        beam_width: int,
        max_tokens: int,
458
    ) -> list[tuple[list[list[int]], list[str]]]:
459
460
461
462
463
464
465
466
467
468
469
470
471
472
        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
473

474
475
    def generate_greedy_logprobs(
        self,
476
        prompts: list[str],
477
        max_tokens: int,
478
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
479
        videos: Optional[PromptVideoInput] = None,
480
        audios: Optional[PromptAudioInput] = None,
481
        **kwargs: Any,
482
    ) -> list[list[torch.Tensor]]:
483
484
485
486
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
487

488
        all_logprobs: list[list[torch.Tensor]] = []
489
        for inputs in all_inputs:
490
            output = self.model.generate(
491
                **self.wrap_device(inputs, device=self.model.device.type),
492
493
494
495
496
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
497
                **kwargs,
498
            )
499
500
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
501
502
503
            all_logprobs.append(seq_logprobs)
        return all_logprobs

504
    def _hidden_states_to_seq_logprobs(
505
        self,
506
507
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
508
509
        output_embeddings = self.model.get_output_embeddings()

510
        seq_logprobs: list[torch.Tensor] = []
511
512
513
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
514
515
                last_hidden_states.to(output_embeddings.weight.device),
                output_embeddings.weight.t(),
516
            )
517
518
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
519
520
521
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

522
523
524
525
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
526
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
527
        num_logprobs: int,
528
    ) -> tuple[list[dict[int, float]], int]:
529
530
531
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

532
        # convert to dict
533
        seq_logprobs_lst: list[dict[int, float]] = []
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
        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,
        )

551
552
    def generate_greedy_logprobs_limit(
        self,
553
        prompts: list[str],
554
555
        max_tokens: int,
        num_logprobs: int,
556
557
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
558
        videos: Optional[PromptVideoInput] = None,
559
        **kwargs: Any,
560
    ) -> list[TokensTextLogprobs]:
561
562
563
564
565
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

566
567
568
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
569

570
        for inputs in all_inputs:
571
            output = self.model.generate(
572
                **self.wrap_device(inputs, device=self.model.device.type),
573
574
575
576
577
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
578
                **kwargs,
579
580
            )

581
582
583
584
585
586
587
588
589
590
591
592
            (
                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))
593

594
595
596
597
598
599
        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,
600
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
601
602
        max_tokens: int,
        num_logprobs: int,
603
        images: Optional[PromptImageInput] = None,
604
        **kwargs: Any,
605
    ) -> list[TokensTextLogprobs]:
606
607
608
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
609

610
611
612
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
613

614
615
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
616
            processor_kwargs: dict[str, Any] = {
617
618
619
620
621
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
622

623
624
            encoder_inputs = self.wrap_device(
                self.processor(**processor_kwargs),
625
626
627
628
629
630
631
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
632
                    self.tokenizer(decoder_prompt,
633
634
635
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
636
637
638
639
640
641
642
643

            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,
644
                **encoder_inputs,
645
646
647
648
649
650
651
652
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
653
654
655
656
657
658
659
660
661
662
663

            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]

664
    def encode(self, prompts: list[str]) -> list[list[torch.Tensor]]:
665
666
        return self.model.encode(prompts)

667
    def predict(self, prompts: list[list[str]]) -> torch.Tensor:
668
669
        return self.model.predict(prompts, convert_to_tensor=True)

670
671
672
673
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
674
        del self.model
675
        cleanup_dist_env_and_memory()
676

Woosuk Kwon's avatar
Woosuk Kwon committed
677

Cyrus Leung's avatar
Cyrus Leung committed
678
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
679
680
681
682
683
684
685
686
687
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
688
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
689
        tokenizer_name: Optional[str] = None,
690
        tokenizer_mode: str = "auto",
691
692
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
693
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
694
        dtype: str = "half",
695
        disable_log_stats: bool = True,
696
        tensor_parallel_size: int = 1,
697
698
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
699
        swap_space: int = 4,
700
        enforce_eager: Optional[bool] = False,
701
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
702
703
704
    ) -> None:
        self.model = LLM(
            model=model_name,
705
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
706
            tokenizer=tokenizer_name,
707
            tokenizer_mode=tokenizer_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
708
709
            trust_remote_code=True,
            dtype=dtype,
710
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
711
            enforce_eager=enforce_eager,
712
            disable_log_stats=disable_log_stats,
713
            tensor_parallel_size=tensor_parallel_size,
714
            max_model_len=max_model_len,
715
716
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
717
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
718
719
        )

720
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
721
        self,
722
        prompts: list[str],
723
        images: Optional[PromptImageInput] = None,
724
725
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
726
    ) -> list[TextPrompt]:
727
        if images is not None:
728
            assert len(prompts) == len(images)
729

730
731
732
733
734
735
        if videos is not None:
            assert len(prompts) == len(videos)

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

736
737
738
        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
739
740
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
741

742
743
        if videos is not None:
            for i, video in enumerate(videos):
Cyrus Leung's avatar
Cyrus Leung committed
744
745
                if video is not None:
                    inputs[i]["multi_modal_data"] = {"video": video}
746
747
748

        if audios is not None:
            for i, audio in enumerate(audios):
Cyrus Leung's avatar
Cyrus Leung committed
749
750
                if audio is not None:
                    inputs[i]["multi_modal_data"] = {"audio": audio}
751
752
753
754
755

        return inputs

    def generate(
        self,
756
        prompts: list[str],
757
758
759
760
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
761
        **kwargs: Any,
762
    ) -> list[tuple[list[list[int]], list[str]]]:
763
764
765
766
767
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

768
        req_outputs = self.model.generate(inputs,
769
770
                                          sampling_params=sampling_params,
                                          **kwargs)
771

772
        outputs: list[tuple[list[list[int]], list[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
773
774
775
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
776
777
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
778
779
            for sample in req_output.outputs:
                output_str = sample.text
780
                output_ids = list(sample.token_ids)
781
782
783
                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
784
785
        return outputs

786
    @staticmethod
787
    def _final_steps_generate_w_logprobs(
788
789
790
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
791
        for req_output in req_outputs:
792
            assert len(req_output.outputs) > 0
793
794
            for sample in req_output.outputs:
                output_str = sample.text
795
                output_ids = list(sample.token_ids)
796
                output_logprobs = sample.logprobs
797
798
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
799
800
        return outputs

801
802
    def generate_w_logprobs(
        self,
803
        prompts: list[str],
804
        sampling_params: SamplingParams,
805
806
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
807
        videos: Optional[PromptVideoInput] = None,
808
        **kwargs: Any,
809
810
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
811
812
813
814
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
815

816
        req_outputs = self.model.generate(inputs,
817
818
                                          sampling_params=sampling_params,
                                          **kwargs)
819
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
827
828

    def generate_encoder_decoder_w_logprobs(
        self,
829
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
830
        sampling_params: SamplingParams,
831
832
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
833
834
835
836
837
838
839
        '''
        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)
840
841
842
843
844
845
        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)
846

Woosuk Kwon's avatar
Woosuk Kwon committed
847
848
    def generate_greedy(
        self,
849
        prompts: list[str],
Woosuk Kwon's avatar
Woosuk Kwon committed
850
        max_tokens: int,
851
        images: Optional[PromptImageInput] = None,
852
853
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
854
        **kwargs: Any,
855
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
856
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
857
858
859
860
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
861
862
                                audios=audios,
                                **kwargs)
863
864
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
865

866
867
    def generate_greedy_logprobs(
        self,
868
        prompts: list[str],
869
870
        max_tokens: int,
        num_logprobs: int,
871
        num_prompt_logprobs: Optional[int] = None,
872
873
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
874
        videos: Optional[PromptVideoInput] = None,
875
876
        stop_token_ids: Optional[list[int]] = None,
        stop: Optional[list[str]] = None,
877
        **kwargs: Any,
878
879
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
880
881
882
883
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
884
            prompt_logprobs=num_prompt_logprobs,
885
886
            stop_token_ids=stop_token_ids,
            stop=stop)
887
888
889
890
891

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
892
893
                                        videos=videos,
                                        **kwargs)
894

895
896
    def generate_encoder_decoder_greedy_logprobs(
        self,
897
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
898
899
        max_tokens: int,
        num_logprobs: int,
900
        num_prompt_logprobs: Optional[int] = None,
901
902
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
903
904
905
906
907
908
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
        )
909
910
911
912
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

913
        return self.generate_encoder_decoder_w_logprobs(
914
915
            encoder_decoder_prompts, greedy_logprobs_params)

916
    def generate_beam_search(
917
        self,
918
        prompts: Union[list[str], list[list[int]]],
919
920
        beam_width: int,
        max_tokens: int,
921
    ) -> list[tuple[list[list[int]], list[str]]]:
youkaichao's avatar
youkaichao committed
922
923
924
925
926
927
        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
            ]
928
929
930
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
931
932
933
934
935
936
937
        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

938
    def classify(self, prompts: list[str]) -> list[list[float]]:
939
940
941
        req_outputs = self.model.classify(prompts)
        return [req_output.outputs.probs for req_output in req_outputs]

Cyrus Leung's avatar
Cyrus Leung committed
942
943
    def encode(
        self,
944
        prompts: list[str],
Cyrus Leung's avatar
Cyrus Leung committed
945
946
947
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
948
    ) -> list[list[float]]:
Cyrus Leung's avatar
Cyrus Leung committed
949
950
951
952
953
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

954
        req_outputs = self.model.embed(inputs)
Cyrus Leung's avatar
Cyrus Leung committed
955
        return [req_output.outputs.embedding for req_output in req_outputs]
956

957
958
    def score(
        self,
959
960
961
        text_1: Union[str, list[str]],
        text_2: Union[str, list[str]],
    ) -> list[float]:
962
        req_outputs = self.model.score(text_1, text_2)
963
        return [req_output.outputs.score for req_output in req_outputs]
964

965
966
967
968
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
        executor = self.model.llm_engine.model_executor
        return executor.apply_model(func)

969
970
971
972
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
973
        del self.model
974
        cleanup_dist_env_and_memory()
975

Woosuk Kwon's avatar
Woosuk Kwon committed
976

977
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
978
979
def vllm_runner():
    return VllmRunner
980
981
982
983
984
985
986
987
988


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={})
989
990
991
992
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
993
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009


@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
1010
1011
1012
1013
1014
1015
1016


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

1017
    return cuda_device_count_stateless()
1018
1019
1020


temp_dir = tempfile.gettempdir()
1021
1022
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1023
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1024
1025
1026
1027


@pytest.fixture
def dummy_opt_path():
1028
1029
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1030
        snapshot_download(repo_id="facebook/opt-125m",
1031
                          local_dir=_dummy_opt_path,
1032
1033
1034
1035
1036
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1037
        with open(json_path) as f:
1038
1039
1040
1041
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
    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)
1056
        with open(json_path) as f:
1057
1058
1059
1060
1061
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074


@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)
1075
        with open(json_path) as f:
1076
1077
1078
1079
1080
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099


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