conftest.py 42.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12

# ruff: noqa

from tblib import pickling_support

# Install support for pickling exceptions so that we can nicely propagate
# failures from tests running in a subprocess.
# This should be run before any custom exception subclasses are defined.
pickling_support.install()

13
import http.server
14
import json
15
import math
16
import mimetypes
17
import os
18
import socket
19
import tempfile
20
21
import threading
from collections.abc import Generator
22
from enum import Enum
23
from typing import Any, Callable, Optional, TypedDict, TypeVar, Union, cast
Woosuk Kwon's avatar
Woosuk Kwon committed
24

25
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
26
27
import pytest
import torch
28
import torch.nn as nn
29
import torch.nn.functional as F
30
from huggingface_hub import snapshot_download
31
from PIL import Image
32
33
from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
                          BatchEncoding, BatchFeature)
34
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
35

36
37
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
38
from vllm import LLM, SamplingParams
39
from vllm.assets.audio import AudioAsset
40
from vllm.assets.image import ImageAsset
41
from vllm.assets.video import VideoAsset
42
from vllm.config import ConvertOption, RunnerOption, _get_and_verify_dtype
43
from vllm.connections import global_http_connection
44
from vllm.distributed import (cleanup_dist_env_and_memory,
45
46
                              init_distributed_environment,
                              initialize_model_parallel)
47
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
48
                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
49
from vllm.logger import init_logger
50
from vllm.multimodal.utils import fetch_image
51
from vllm.outputs import RequestOutput
52
from vllm.sampling_params import BeamSearchParams
53
from vllm.sequence import Logprob
54
from vllm.transformers_utils.utils import maybe_model_redirect
55

56
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
57

58
59
60
_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")]
61
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
62

Cyrus Leung's avatar
Cyrus Leung committed
63
_M = TypeVar("_M")
64

65
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
66
67

PromptImageInput = _PromptMultiModalInput[Image.Image]
68
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
69
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
70

71

72
def _read_prompts(filename: str) -> list[str]:
73
    with open(filename) as f:
74
75
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
76
77


78
class ImageAssetPrompts(TypedDict):
79
80
    stop_sign: str
    cherry_blossom: str
81
82


83
class ImageTestAssets(list[ImageAsset]):
84
85

    def __init__(self) -> None:
86
87
88
89
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
90

91
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
92
93
94
95
96
97
        """
        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.
        """
98
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
99
100


101
102
class VideoAssetPrompts(TypedDict):
    baby_reading: str
103
104


105
class VideoTestAssets(list[VideoAsset]):
106
107
108

    def __init__(self) -> None:
        super().__init__([
109
            VideoAsset("baby_reading"),
110
111
        ])

112
113
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
114
115


116
class AudioAssetPrompts(TypedDict):
117
118
119
120
    mary_had_lamb: str
    winning_call: str


121
class AudioTestAssets(list[AudioAsset]):
122
123
124
125
126
127
128

    def __init__(self) -> None:
        super().__init__([
            AudioAsset("mary_had_lamb"),
            AudioAsset("winning_call"),
        ])

129
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
130
131
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

132

133
IMAGE_ASSETS = ImageTestAssets()
134
"""Singleton instance of {class}`ImageTestAssets`."""
135
VIDEO_ASSETS = VideoTestAssets()
136
"""Singleton instance of {class}`VideoTestAssets`."""
137
AUDIO_ASSETS = AudioTestAssets()
138
"""Singleton instance of {class}`AudioTestAssets`."""
139
140


141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@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
161
@pytest.fixture(params=[True, False])
162
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
163
164
165
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
166
    skip_v0 = request.node.get_closest_marker("skip_v0")
Joe Runde's avatar
Joe Runde committed
167
168
169
170
171
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
172
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
173
    else:
174
175
        if skip_v0:
            pytest.skip("Skipping test on vllm V0")
176
177
178
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
179
180


181
182
183
184
185
186
187
@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


188
189
190
191
192
193
194
195
196
197
198
199
@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
200
    cleanup_dist_env_and_memory()
201
202


203
@pytest.fixture()
204
def should_do_global_cleanup_after_test(request) -> bool:
205
206
207
208
    """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.
    """
209

210
    return not request.node.get_closest_marker("skip_global_cleanup")
211
212


213
@pytest.fixture(autouse=True)
214
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
215
    yield
216
    if should_do_global_cleanup_after_test:
217
        cleanup_dist_env_and_memory()
218
219


220
221
222
223
224
225
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


234
235
236
237
238
239
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


240
241
242
243
244
245
246
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


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


255
@pytest.fixture(scope="session")
256
def image_assets() -> ImageTestAssets:
257
258
259
    return IMAGE_ASSETS


260
@pytest.fixture(scope="session")
261
def video_assets() -> VideoTestAssets:
262
263
264
    return VIDEO_ASSETS


265
@pytest.fixture(scope="session")
266
def audio_assets() -> AudioTestAssets:
267
268
269
    return AUDIO_ASSETS


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

Woosuk Kwon's avatar
Woosuk Kwon committed
273
274
275

class HfRunner:

276
    def get_default_device(self):
277
        from vllm.platforms import current_platform
278

279
280
        return ("cpu"
                if current_platform.is_cpu() else current_platform.device_type)
281
282

    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
283
284
285
        if x is None or isinstance(x, (bool, )):
            return x

286
        if device is None:
287
            device = self.device
288

289
290
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
291

292
293
294
295
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
296

Woosuk Kwon's avatar
Woosuk Kwon committed
297
298
299
    def __init__(
        self,
        model_name: str,
300
        dtype: str = "auto",
301
        *,
302
        model_kwargs: Optional[dict[str, Any]] = None,
303
        trust_remote_code: bool = True,
304
        is_sentence_transformer: bool = False,
305
        is_cross_encoder: bool = False,
306
        skip_tokenizer_init: bool = False,
307
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
Woosuk Kwon's avatar
Woosuk Kwon committed
308
    ) -> None:
309
        model_name = maybe_model_redirect(model_name)
310
        self.model_name = model_name
311

312
313
        self.config = AutoConfig.from_pretrained(
            model_name,
314
            trust_remote_code=trust_remote_code,
315
316
        )
        self.device = self.get_default_device()
317
318
319
320
321
322
        self.dtype = torch_dtype = _get_and_verify_dtype(
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
323
324
325
326

        model_kwargs = model_kwargs if model_kwargs is not None else {}
        model_kwargs.setdefault("torch_dtype", torch_dtype)

327
        if is_sentence_transformer:
328
329
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
330
331
332
333
334

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
335
                trust_remote_code=trust_remote_code,
336
            )
337
338
339
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
340
341
342
343
344

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
345
                trust_remote_code=trust_remote_code,
346
            )
347
        else:
348
349
            model = auto_cls.from_pretrained(
                model_name,
350
                trust_remote_code=trust_remote_code,
351
352
353
                **model_kwargs,
            )

354
355
356
357
358
359
            # in case some unquantized custom models are not in same dtype
            if (getattr(model, "quantization_method", None) is None
                    and any(p.dtype != self.dtype
                            for p in model.parameters())):
                model = model.to(dtype=self.dtype)

360
361
362
            if (getattr(model, "quantization_method", None) != "bitsandbytes"
                    and len({p.device
                             for p in model.parameters()}) < 2):
363
                model = model.to(device=self.device)
364
365

            self.model = model
366

367
368
369
370
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
371
                trust_remote_code=trust_remote_code,
372
            )
373

374
375
376
377
378
379
        # 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,
380
            trust_remote_code=trust_remote_code,
381
        )
382
383
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
384

385
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
386
        self,
387
        prompts: list[str],
388
        images: Optional[PromptImageInput] = None,
389
390
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
391
    ) -> list[Union[BatchFeature, BatchEncoding]]:
392
        if images is not None:
393
            assert len(prompts) == len(images)
394

395
396
397
398
399
400
        if videos is not None:
            assert len(prompts) == len(videos)

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

401
        all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
402
        for i, prompt in enumerate(prompts):
403
            processor_kwargs: dict[str, Any] = {
404
405
406
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
407
408
409
410
            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
411
412
413
414
415
416
417
418
419
            if audios is not None and (audio_inputs := audios[i]) is not None:
                # HACK - not all processors take sampling_rate; we should
                # clean this up in the future.
                if len(audio_inputs) == 2:
                    audio, sr = audio_inputs
                    processor_kwargs["audio"] = audio
                    processor_kwargs["sampling_rate"] = sr
                else:
                    processor_kwargs["audio"] = audio_inputs
420
421

            inputs = self.processor(**processor_kwargs)
422
423
            if isinstance(inputs, BatchFeature):
                inputs = inputs.to(dtype=self.dtype)
424

425
426
427
428
            all_inputs.append(inputs)

        return all_inputs

429
430
431
432
433
434
435
436
437
    def get_prompt_embeddings(self, prompts: list[str]) -> list[torch.Tensor]:
        all_inputs = self.get_inputs(prompts)
        embeddings = []
        for inputs in all_inputs:
            input_ids = self.wrap_device(inputs)["input_ids"]
            embedding = self.model.get_input_embeddings()(input_ids).squeeze(0)
            embeddings.append(embedding)
        return embeddings

438
    def classify(self, prompts: list[str]) -> list[str]:
439
440
441
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
442
443
        problem_type = getattr(self.config, "problem_type", "")

444
445
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
446
447
448
449
450
451
            if problem_type == "regression":
                logits = output.logits[0].tolist()
            elif problem_type == "multi_label_classification":
                logits = output.logits.sigmoid()[0].tolist()
            else:
                logits = output.logits.softmax(dim=-1)[0].tolist()
452
453
454
455
            outputs.append(logits)

        return outputs

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

469
        outputs: list[tuple[list[list[int]], list[str]]] = []
470
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
471
            output_ids = self.model.generate(
472
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
473
474
475
                use_cache=True,
                **kwargs,
            )
476
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
477
478
479
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
480
481
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
482
483
484
485
486
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
487
        prompts: list[str],
Woosuk Kwon's avatar
Woosuk Kwon committed
488
        max_tokens: int,
489
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
490
        videos: Optional[PromptVideoInput] = None,
491
        audios: Optional[PromptAudioInput] = None,
492
        **kwargs: Any,
493
    ) -> list[tuple[list[int], str]]:
494
495
        outputs = self.generate(prompts,
                                do_sample=False,
496
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
497
                                images=images,
498
499
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
500
                                **kwargs)
501
502
503

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
504
505
506

    def generate_beam_search(
        self,
507
        prompts: list[str],
508
509
        beam_width: int,
        max_tokens: int,
510
511
512
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
513
    ) -> list[tuple[list[list[int]], list[str]]]:
514
515
516
517
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
518
519
520
521
522
                                num_return_sequences=beam_width,
                                images=images,
                                videos=videos,
                                audios=audios)

523
524
525
526
527
528
529
530
531
        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
532

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

547
        all_logprobs: list[list[torch.Tensor]] = []
548
        for inputs in all_inputs:
549
            output = self.model.generate(
550
                **self.wrap_device(inputs),
551
552
553
554
555
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
556
                **kwargs,
557
            )
558
559
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
560
561
562
            all_logprobs.append(seq_logprobs)
        return all_logprobs

563
    def _hidden_states_to_seq_logprobs(
564
        self,
565
566
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
567
568
        output_embeddings = self.model.get_output_embeddings()

569
        seq_logprobs: list[torch.Tensor] = []
570
571
572
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
573
574
575
576
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
577
                output_embeddings.weight.t(),
578
            )
579
580
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
581
582
583
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

584
585
586
587
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
588
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
589
        num_logprobs: Optional[int],
590
    ) -> tuple[list[dict[int, float]], int]:
591
592
593
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

594
        # convert to dict
595
        seq_logprobs_lst: list[dict[int, float]] = []
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
        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,
        )

613
614
    def generate_greedy_logprobs_limit(
        self,
615
        prompts: list[str],
616
        max_tokens: int,
617
        num_logprobs: Optional[int],
618
619
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
620
        videos: Optional[PromptVideoInput] = None,
621
        **kwargs: Any,
622
    ) -> list[TokensTextLogprobs]:
623
624
625
626
627
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

628
629
630
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
631

632
        for inputs in all_inputs:
633
            output = self.model.generate(
634
                **self.wrap_device(inputs),
635
636
637
638
639
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
640
                **kwargs,
641
642
            )

643
644
645
646
647
648
649
650
651
652
653
654
            (
                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))
655

656
657
658
659
        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]

660
661
662
    def encode(self, prompts: list[str], *args,
               **kwargs) -> list[list[torch.Tensor]]:
        return self.model.encode(prompts, *args, **kwargs)
663

664
665
666
667
668
669
    def predict(self, prompts: list[list[str]], *args,
                **kwargs) -> torch.Tensor:
        return self.model.predict(prompts,
                                  *args,
                                  convert_to_tensor=True,
                                  **kwargs)
670

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

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

Woosuk Kwon's avatar
Woosuk Kwon committed
678

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


class VllmRunner:
685
686
    """
    The default value of some arguments have been modified from
687
    {class}`~vllm.LLM` as follows:
688

689
690
691
    - `trust_remote_code`: Set to `True` instead of `False` for convenience.
    - `seed`: Set to `0` instead of `None` for test reproducibility.
    - `max_model_len`: Set to `1024` instead of `None` to reduce memory usage.
692
693
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
694
695
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
696
    - `enforce_eager`: Set to `False` to test CUDA graph.
697
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
698
699
700
701

    def __init__(
        self,
        model_name: str,
702
703
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
704
        tokenizer_name: Optional[str] = None,
705
        tokenizer_mode: str = "auto",
706
707
        trust_remote_code: bool = True,
        seed: Optional[int] = 0,
708
        max_model_len: Optional[int] = 1024,
709
        dtype: str = "auto",
710
        disable_log_stats: bool = True,
711
        tensor_parallel_size: int = 1,
712
        block_size: int = 16 if not torch.xpu.is_available() else 64,
713
        enable_chunked_prefill: Optional[bool] = False,
714
        swap_space: int = 4,
715
        enforce_eager: Optional[bool] = False,
716
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
717
    ) -> None:
718
        self.llm = LLM(
Woosuk Kwon's avatar
Woosuk Kwon committed
719
            model=model_name,
720
721
            runner=runner,
            convert=convert,
Woosuk Kwon's avatar
Woosuk Kwon committed
722
            tokenizer=tokenizer_name,
723
            tokenizer_mode=tokenizer_mode,
724
            trust_remote_code=trust_remote_code,
Woosuk Kwon's avatar
Woosuk Kwon committed
725
            dtype=dtype,
726
            seed=seed,
727
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
728
            enforce_eager=enforce_eager,
729
            disable_log_stats=disable_log_stats,
730
            tensor_parallel_size=tensor_parallel_size,
731
            max_model_len=max_model_len,
732
733
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
734
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
735
736
        )

737
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
738
        self,
739
        prompts: Union[list[str], list[torch.Tensor], list[int]],
740
        images: Optional[PromptImageInput] = None,
741
742
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
743
    ) -> list[TextPrompt]:
744

745
746
747
748
749
        if any(x is not None and len(x) != len(prompts)
               for x in [images, videos, audios]):
            raise ValueError(
                "All non-None multimodal inputs must have the same length as "
                "prompts")
750

751
752
753
754
755
756
757
758
759
760
        inputs = []
        for i, prompt in enumerate(prompts):
            multi_modal_data = {}
            if images is not None and (image := images[i]) is not None:
                multi_modal_data["image"] = image
            if videos is not None and (video := videos[i]) is not None:
                multi_modal_data["video"] = video
            if audios is not None and (audio := audios[i]) is not None:
                multi_modal_data["audio"] = audio

761
            text_prompt_kwargs: dict[str, Any] = {
762
763
                "multi_modal_data": multi_modal_data or None
            }
764
765
766
767
768
769
770
            if isinstance(prompt, str):
                text_prompt_kwargs["prompt"] = prompt
            elif isinstance(prompt, list):
                text_prompt_kwargs["prompt_token_ids"] = prompt
            else:
                text_prompt_kwargs["prompt_embeds"] = prompt

771
            inputs.append(TextPrompt(**text_prompt_kwargs))
772
773
774
775
776

        return inputs

    def generate(
        self,
777
        prompts: Union[list[str], list[torch.Tensor]],
778
779
780
781
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
782
        **kwargs: Any,
783
    ) -> list[tuple[list[list[int]], list[str]]]:
784
785
786
787
788
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

789
790
791
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
792

793
        outputs: list[tuple[list[list[int]], list[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
794
795
796
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
797
798
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
799
800
            for sample in req_output.outputs:
                output_str = sample.text
801
                output_ids = list(sample.token_ids)
802
                req_sample_output_ids.append(prompt_ids + output_ids)
803
                req_sample_output_strs.append((prompt_str or "") + output_str)
804
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
805
806
        return outputs

807
    @staticmethod
808
    def _final_steps_generate_w_logprobs(
809
810
811
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
812
        for req_output in req_outputs:
813
            assert len(req_output.outputs) > 0
814
815
            for sample in req_output.outputs:
                output_str = sample.text
816
                output_ids = list(sample.token_ids)
817
                output_logprobs = sample.logprobs
818
819
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
820
821
        return outputs

822
823
    def generate_w_logprobs(
        self,
824
        prompts: list[str],
825
        sampling_params: SamplingParams,
826
827
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
828
        videos: Optional[PromptVideoInput] = None,
829
        **kwargs: Any,
830
831
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
832
833
834
835
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
836

837
838
839
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
840
841
842
843
844
845
846

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

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

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

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

896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
    def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:
        """
        Return the perplexity score associated with generating the prompts

        :param prompts: list of prompts to score
        :return: perplexity score of each prompt
        """
        outputs = self.generate_greedy_logprobs(prompts,
                                                max_tokens=1,
                                                num_logprobs=None,
                                                num_prompt_logprobs=0)

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
            token_datas = cast(list[Optional[dict[int, Logprob]]], output[3])
            assert token_datas[0] is None
            token_log_probs = []
            for token_data in token_datas[1:]:
                assert token_data is not None
                assert len(token_data) == 1
                token_log_prob = list(token_data.values())[0].logprob
                token_log_probs.append(token_log_prob)

            perplexity = math.exp(-sum(token_log_probs) / len(token_log_probs))
            perplexities.append(perplexity)

        return perplexities

925
    def generate_beam_search(
926
        self,
927
        prompts: list[str],
928
929
        beam_width: int,
        max_tokens: int,
930
931
932
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
933
        concurrency_limit: Optional[int] = None,
934
    ) -> list[tuple[list[list[int]], list[str]]]:
935
936
937
938
939
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

940
941
942
943
        outputs = self.llm.beam_search(inputs,
                                       BeamSearchParams(beam_width=beam_width,
                                                        max_tokens=max_tokens),
                                       concurrency_limit=concurrency_limit)
944
945
946
947
948
949
950
        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

951
    def classify(self, prompts: list[str]) -> list[list[float]]:
952
        req_outputs = self.llm.classify(prompts)
953
954
        return [req_output.outputs.probs for req_output in req_outputs]

955
956
957
958
959
960
961
    def embed(self,
              prompts: list[str],
              images: Optional[PromptImageInput] = None,
              videos: Optional[PromptVideoInput] = None,
              audios: Optional[PromptAudioInput] = None,
              *args,
              **kwargs) -> list[list[float]]:
Cyrus Leung's avatar
Cyrus Leung committed
962
963
964
965
966
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

967
        req_outputs = self.llm.embed(inputs, *args, **kwargs)
Cyrus Leung's avatar
Cyrus Leung committed
968
        return [req_output.outputs.embedding for req_output in req_outputs]
969

970
    def encode(self, prompts: list[str]) -> list[list[float]]:
971
        req_outputs = self.llm.encode(prompts)
972
973
        return [req_output.outputs.data for req_output in req_outputs]

974
975
976
977
    def reward(self, prompts: list[str]) -> list[list[float]]:
        req_outputs = self.llm.reward(prompts)
        return [req_output.outputs.data for req_output in req_outputs]

978
979
    def score(
        self,
980
981
        text_1: Union[str, list[str]],
        text_2: Union[str, list[str]],
982
983
        *args,
        **kwargs,
984
    ) -> list[float]:
985
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
986
        return [req_output.outputs.score for req_output in req_outputs]
987

988
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
989
990
991
992
993
994
995
996
997
998
999
        if hasattr(self.llm.llm_engine, "model_executor"):
            # This works either in V0 or in V1 with
            # VLLM_ENABLE_V1_MULTIPROCESSING=0
            executor = self.llm.llm_engine.model_executor
            return executor.apply_model(func)

        # This works in V1 with VLLM_ALLOW_INSECURE_SERIALIZATION=1
        def _apply_model(self):
            return func(self.get_model())

        return self.llm.llm_engine.collective_rpc(_apply_model)
1000

1001
1002
1003
    def get_llm(self) -> LLM:
        return self.llm

1004
1005
1006
1007
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1008
        del self.llm
1009
        cleanup_dist_env_and_memory()
1010

Woosuk Kwon's avatar
Woosuk Kwon committed
1011

1012
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1013
1014
def vllm_runner():
    return VllmRunner
1015
1016


1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
@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
1031
1032
1033
1034
1035
1036
1037


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

1038
1039
    from vllm.platforms import current_platform
    return current_platform.device_count()
1040
1041
1042


temp_dir = tempfile.gettempdir()
1043
1044
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1045
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1046
1047
1048
1049


@pytest.fixture
def dummy_opt_path():
1050
1051
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1052
        snapshot_download(repo_id="facebook/opt-125m",
1053
                          local_dir=_dummy_opt_path,
1054
1055
1056
1057
1058
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1059
        with open(json_path) as f:
1060
1061
1062
1063
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
    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)
1078
        with open(json_path) as f:
1079
1080
1081
1082
1083
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096


@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)
1097
        with open(json_path) as f:
1098
1099
1100
1101
1102
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120


# 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:
1121
            item.add_marker(skip_optional)
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133


@pytest.fixture(scope="session")
def cli_config_file():
    """Return the path to the CLI config file."""
    return os.path.join(_TEST_DIR, "config", "test_config.yaml")


@pytest.fixture(scope="session")
def cli_config_file_with_model():
    """Return the path to the CLI config file with model."""
    return os.path.join(_TEST_DIR, "config", "test_config_with_model.yaml")
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249


class AssetHandler(http.server.BaseHTTPRequestHandler):
    # _IMAGE_CACHE : Dict[str, bytes] = {}

    def log_message(self, *args, **kwargs):
        pass

    def do_GET(self):
        # Accepts paths like: /1280px-Venn_diagram_rgb.jpg
        filename = self.path.lstrip("/")
        if not filename or "." not in filename:
            self.send_error(404, "Missing filename (expected /<name>.<ext>)")
            return

        base, ext = filename.rsplit(".", 1)
        ext = ext.lower()

        if ext not in ["jpg", "png"]:
            self.send_error(404, f"Unsupported extension: .{ext}")
            return

        try:
            data = ImageAsset(base).read_bytes(ext=ext)
        except Exception as e:
            self.send_error(500, f"Failed to load asset: {ext} {base} {e} ")
            return

        ctype, _ = mimetypes.guess_type(filename)
        if ctype is None:
            ctype = {"jpg": "image/jpg", "png": "image/png"}[ext]
        self.send_response(200)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)


def _find_free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


class LocalAssetServer:

    address: str
    port: int
    server: Optional[http.server.ThreadingHTTPServer]
    thread: Optional[threading.Thread]

    def __init__(self, address: str = "127.0.0.1") -> None:
        self.address = address
        self.port = -1
        self.server = None
        self.thread = None

    def __enter__(self):
        self.port = _find_free_port()
        self.server = http.server.ThreadingHTTPServer(
            (self.address, self.port), AssetHandler)
        self.thread = threading.Thread(target=self.server.serve_forever,
                                       daemon=True)
        self.thread.start()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self.server:
            self.server.shutdown()
            del self.server

        if self.thread:
            self.thread.join()
            del self.thread

        if exc_type is None:
            return None

        return False

    @property
    def base_url(self) -> str:
        assert self.port is not None
        return f"http://{self.address}:{self.port}"

    def url_for(self, name: str) -> str:
        """e.g., name='RGBA_comp.png' -> 'http://127.0.0.1:PORT/RGBA_comp.png'"""
        return f"{self.base_url}/{name}"

    def get_image_asset(self, name: str) -> Image.Image:
        return fetch_image(self.url_for(name))


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port. 
    The server currently servers images at:
    http://127.0.0.1:<port>/<name>.<ext>
    """
    with LocalAssetServer() as srv:
        yield srv


@pytest.fixture
def image_url(request, local_asset_server) -> str:
    # request.param is one of the IMAGE_ASSETS filenames
    name = request.param
    return local_asset_server.url_for(name)


@pytest.fixture
def image_urls(request, local_asset_server) -> list[str]:
    """Indirect fixture: takes a list of names, returns list of full URLs."""
    names: list[str] = request.param
    return [local_asset_server.url_for(name) for name in names]