conftest.py 49.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
import contextlib
import pathlib
from copy import deepcopy
6
7
8

from tblib import pickling_support

9
10
11
# Import fixture
from tests.v1.entrypoints.conftest import sample_json_schema  # noqa

12
13
# ruff: noqa

14
15
16
17
18
# 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()

19
import http.server
20
import json
21
import math
22
import mimetypes
23
import os
24
import socket
25
import tempfile
26
27
import threading
from collections.abc import Generator
28
from contextlib import nullcontext
29
from enum import Enum
30
from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING
Woosuk Kwon's avatar
Woosuk Kwon committed
31

32
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
33
34
import pytest
import torch
35
import torch.nn as nn
36
import torch.nn.functional as F
37
from huggingface_hub import snapshot_download
38
from PIL import Image
39
40
41
42
43
44
45
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BatchEncoding,
    BatchFeature,
)
46
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
47

48
49
50
51
52
from tests.models.utils import (
    TokensTextLogprobs,
    TokensTextLogprobsPromptLogprobs,
    softmax,
)
53
from vllm import LLM, SamplingParams, envs
54
from vllm.assets.audio import AudioAsset
55
from vllm.assets.image import ImageAsset
56
from vllm.assets.video import VideoAsset
57
from vllm.config.model import ConvertOption, RunnerOption, _get_and_verify_dtype
58
from vllm.connections import global_http_connection
59
60
61
62
63
from vllm.distributed import (
    cleanup_dist_env_and_memory,
    init_distributed_environment,
    initialize_model_parallel,
)
64
from vllm.logger import init_logger
65
from vllm.logprobs import Logprob
66
from vllm.multimodal.base import MediaWithBytes
67
from vllm.multimodal.utils import fetch_image
68
from vllm.outputs import RequestOutput
69
from vllm.sampling_params import BeamSearchParams
70
from vllm.transformers_utils.utils import maybe_model_redirect
71
from vllm.utils.collection_utils import is_list_of
72
from vllm.utils.torch_utils import set_default_torch_num_threads
73

74
75
76
from torch._inductor.utils import fresh_cache


77
78
79
80
81
if TYPE_CHECKING:
    from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
    from transformers.generation.utils import GenerateOutput


82
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
83

84
85
86
_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")]
87
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
88

Cyrus Leung's avatar
Cyrus Leung committed
89
_M = TypeVar("_M")
90

91
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
92
93

PromptImageInput = _PromptMultiModalInput[Image.Image]
94
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
95
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
96

97

98
def _read_prompts(filename: str) -> list[str]:
99
    with open(filename) as f:
100
101
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
102
103


104
class ImageAssetPrompts(TypedDict):
105
106
    stop_sign: str
    cherry_blossom: str
107
108


109
class ImageTestAssets(list[ImageAsset]):
110
    def __init__(self) -> None:
111
112
113
114
115
116
        super().__init__(
            [
                ImageAsset("stop_sign"),
                ImageAsset("cherry_blossom"),
            ]
        )
117

118
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
119
120
121
122
123
124
        """
        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.
        """
125
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
126
127


128
129
class VideoAssetPrompts(TypedDict):
    baby_reading: str
130
131


132
class VideoTestAssets(list[VideoAsset]):
133
    def __init__(self) -> None:
134
135
136
137
138
        super().__init__(
            [
                VideoAsset("baby_reading"),
            ]
        )
139

140
141
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
142
143


144
class AudioAssetPrompts(TypedDict):
145
146
147
148
    mary_had_lamb: str
    winning_call: str


149
class AudioTestAssets(list[AudioAsset]):
150
    def __init__(self) -> None:
151
152
153
154
155
156
        super().__init__(
            [
                AudioAsset("mary_had_lamb"),
                AudioAsset("winning_call"),
            ]
        )
157

158
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
159
160
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

161

162
IMAGE_ASSETS = ImageTestAssets()
163
"""Singleton instance of {class}`ImageTestAssets`."""
164
VIDEO_ASSETS = VideoTestAssets()
165
"""Singleton instance of {class}`VideoTestAssets`."""
166
AUDIO_ASSETS = AudioTestAssets()
167
"""Singleton instance of {class}`AudioTestAssets`."""
168
169


170
171
172
173
174
175
176
@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


177
178
179
180
181
182
183
184
185
186
187
188
@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
189
    cleanup_dist_env_and_memory()
190
191


192
@pytest.fixture()
193
def should_do_global_cleanup_after_test(request) -> bool:
194
195
196
197
    """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.
    """
198

199
    return not request.node.get_closest_marker("skip_global_cleanup")
200
201


202
@pytest.fixture(autouse=True)
203
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
204
    yield
205
    if should_do_global_cleanup_after_test:
206
        cleanup_dist_env_and_memory()
207
208


209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
@pytest.fixture
def workspace_init():
    """Initialize the workspace manager for tests that need it.

    This fixture initializes the workspace manager with a CUDA device
    if available, and resets it after the test completes. Tests that
    create a full vLLM engine should NOT use this fixture as the engine
    will initialize the workspace manager itself.
    """
    from vllm.v1.worker.workspace import (
        init_workspace_manager,
        reset_workspace_manager,
    )

    if torch.cuda.is_available():
        device = torch.device("cuda:0")
        init_workspace_manager(device)
    yield
    reset_workspace_manager()


230
231
232
233
234
235
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
236
@pytest.fixture
237
def example_prompts() -> list[str]:
238
    return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
239
240


241
242
243
244
245
246
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


247
248
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
249

250
251
252
253
254
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


255
@pytest.fixture
256
def example_long_prompts() -> list[str]:
257
    return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
Woosuk Kwon's avatar
Woosuk Kwon committed
258
259


260
@pytest.fixture(scope="session")
261
def image_assets() -> ImageTestAssets:
262
263
264
    return IMAGE_ASSETS


265
@pytest.fixture(scope="session")
266
def video_assets() -> VideoTestAssets:
267
268
269
    return VIDEO_ASSETS


270
@pytest.fixture(scope="session")
271
def audio_assets() -> AudioTestAssets:
272
273
274
    return AUDIO_ASSETS


275
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
276
_R = TypeVar("_R")
277

Woosuk Kwon's avatar
Woosuk Kwon committed
278
279

class HfRunner:
280
    def get_default_device(self):
281
        from vllm.platforms import current_platform
282

283
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
284

285
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
286
        if x is None or isinstance(x, (bool,)):
287
288
            return x

289
        if device is None:
290
            device = self.device
291

292
293
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
294

295
296
297
298
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
299

Woosuk Kwon's avatar
Woosuk Kwon committed
300
301
302
    def __init__(
        self,
        model_name: str,
303
        dtype: str = "auto",
304
        *,
305
        model_kwargs: dict[str, Any] | None = None,
306
        trust_remote_code: bool = True,
307
        is_sentence_transformer: bool = False,
308
        is_cross_encoder: bool = False,
309
        skip_tokenizer_init: bool = False,
310
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
311
        # Set this to avoid hanging issue
312
        default_torch_num_threads: int | None = None,
313
    ) -> None:
314
315
316
317
318
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336

        with init_ctx:
            self._init(
                model_name=model_name,
                dtype=dtype,
                model_kwargs=model_kwargs,
                trust_remote_code=trust_remote_code,
                is_sentence_transformer=is_sentence_transformer,
                is_cross_encoder=is_cross_encoder,
                skip_tokenizer_init=skip_tokenizer_init,
                auto_cls=auto_cls,
            )

    def _init(
        self,
        model_name: str,
        dtype: str = "auto",
        *,
337
        model_kwargs: dict[str, Any] | None = None,
338
339
340
341
342
        trust_remote_code: bool = True,
        is_sentence_transformer: bool = False,
        is_cross_encoder: bool = False,
        skip_tokenizer_init: bool = False,
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
Woosuk Kwon's avatar
Woosuk Kwon committed
343
    ) -> None:
344
        model_name = maybe_model_redirect(model_name)
345
        self.model_name = model_name
346

347
348
        self.config = AutoConfig.from_pretrained(
            model_name,
349
            trust_remote_code=trust_remote_code,
350
351
        )
        self.device = self.get_default_device()
352
        self.dtype = dtype = _get_and_verify_dtype(
353
354
355
356
357
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
358
359

        model_kwargs = model_kwargs if model_kwargs is not None else {}
360
        model_kwargs.setdefault("dtype", dtype)
361

362
        if is_sentence_transformer:
363
364
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
365
366
367
368
369

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
370
                trust_remote_code=trust_remote_code,
371
            )
372
373
374
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
375
376
377
378
379

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
380
                trust_remote_code=trust_remote_code,
381
            )
382
        else:
383
384
385
386
387
388
389
            model = cast(
                nn.Module,
                auto_cls.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                    **model_kwargs,
                ),
390
391
            )

392
            # in case some unquantized custom models are not in same dtype
393
394
395
            if getattr(model, "quantization_method", None) is None and any(
                p.dtype != self.dtype for p in model.parameters()
            ):
396
397
                model = model.to(dtype=self.dtype)

398
399
400
401
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
402
                model = model.to(device=self.device)
403
404

            self.model = model
405

406
        if not skip_tokenizer_init:
407
408
409
410
411
412
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
                    dtype=dtype,
                    trust_remote_code=trust_remote_code,
                )
413
            )
414

415
416
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
417
        from transformers import AutoProcessor
418

419
420
        self.processor = AutoProcessor.from_pretrained(
            model_name,
421
            dtype=dtype,
422
            trust_remote_code=trust_remote_code,
423
        )
424
425
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
426

427
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
428
        self,
429
430
431
432
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
433
        tokenization_kwargs: dict[str, Any] | None = None,
434
    ) -> list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]]:
435
        if images is not None:
436
            assert len(prompts) == len(images)
437

438
439
440
441
442
443
        if videos is not None:
            assert len(prompts) == len(videos)

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

444
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
445
        for i, prompt in enumerate(prompts):
446
            if isinstance(prompt, str):
447
448
449
450
451
452
453
454
455
456
457
458
                # Create a copy to avoid modifying the original dict
                processor_kwargs = (
                    tokenization_kwargs.copy()
                    if tokenization_kwargs is not None
                    else {}
                )
                processor_kwargs.update(
                    {
                        "text": prompt,
                        "return_tensors": "pt",
                    }
                )
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
                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_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

                inputs = self.processor(**processor_kwargs)
                if isinstance(inputs, BatchFeature):
                    inputs = inputs.to(dtype=self.dtype)
                all_inputs.append(inputs)
            else:
                # check that prompt is (batched) list of integers (token ids)
                if not is_list_of(prompt, typ=int, check="all"):
                    raise ValueError(
                        "Prompt must be a list of ints corresponding to the prompt token ids."
                    )
                # check that no multimodal input is provided
                if images or videos or audios:
                    raise ValueError(
                        "When providing prompt token ids multimodal inputs are not supported."
                    )
                input_dict = {
                    "input_ids": torch.tensor(prompt, dtype=torch.long).unsqueeze(0),
                }
                all_inputs.append(input_dict)
492
493
494

        return all_inputs

495
496
497
498
499
500
501
502
503
    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

504
    def classify(self, prompts: list[str]) -> list[list[float]]:
505
506
        # output is final logits
        all_inputs = self.get_inputs(prompts)
507
        outputs: list[list[float]] = []
508
509
        problem_type = getattr(self.config, "problem_type", "")

510
511
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
512
513
514

            assert isinstance(output.logits, torch.Tensor)

515
516
517
518
519
            if problem_type == "regression":
                logits = output.logits[0].tolist()
            elif problem_type == "multi_label_classification":
                logits = output.logits.sigmoid()[0].tolist()
            else:
520
                logits = softmax(output.logits)[0].tolist()
521
522
523
524
            outputs.append(logits)

        return outputs

525
526
    def generate(
        self,
527
528
529
530
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
531
        **kwargs: Any,
532
    ) -> list[tuple[list[list[int]], list[str]]]:
533
534
535
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
536

537
        outputs: list[tuple[list[list[int]], list[str]]] = []
538
        for inputs in all_inputs:
539
            output_ids: torch.Tensor = self.model.generate(
540
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
541
542
543
                use_cache=True,
                **kwargs,
            )
544
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
545
546
547
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
548
            )
549
            outputs.append((output_ids.cpu().tolist(), output_str))
Woosuk Kwon's avatar
Woosuk Kwon committed
550
551
552
553
        return outputs

    def generate_greedy(
        self,
554
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
555
        max_tokens: int,
556
557
558
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
559
        **kwargs: Any,
560
    ) -> list[tuple[list[int], str]]:
561
562
563
564
565
566
567
568
569
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
570

571
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
572
573
574

    def generate_beam_search(
        self,
575
        prompts: list[str],
576
577
        beam_width: int,
        max_tokens: int,
578
579
580
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
581
    ) -> list[tuple[list[list[int]], list[str]]]:
582
583
584
585
586
587
588
589
590
591
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            num_beams=beam_width,
            num_return_sequences=beam_width,
            images=images,
            videos=videos,
            audios=audios,
        )
592

593
594
595
596
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
597
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
598
599
600
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
601

602
603
    def generate_greedy_logprobs(
        self,
604
        prompts: list[str],
605
        max_tokens: int,
606
607
608
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
609
        **kwargs: Any,
610
    ) -> list[list[torch.Tensor]]:
611
612
613
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
614

615
        all_logprobs: list[list[torch.Tensor]] = []
616
        for inputs in all_inputs:
617
            output: "GenerateOutput" = self.model.generate(
618
                **self.wrap_device(inputs),
619
620
621
622
623
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
624
                **kwargs,
625
            )
626
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
627
628
629
            all_logprobs.append(seq_logprobs)
        return all_logprobs

630
    def _hidden_states_to_seq_logprobs(
631
        self,
632
633
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
634
635
        output_embeddings = self.model.get_output_embeddings()

636
        seq_logprobs: list[torch.Tensor] = []
637
638
639
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
640
641
642
643
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
644
                output_embeddings.weight.t(),
645
            )
646
647
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
648
649
650
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

651
652
653
654
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
655
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
656
        num_logprobs: int | None,
657
    ) -> tuple[list[dict[int, float]], int]:
658
659
660
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

661
        # convert to dict
662
        seq_logprobs_lst: list[dict[int, float]] = []
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
        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,
        )

680
681
    def generate_greedy_logprobs_limit(
        self,
682
        prompts: list[str],
683
        max_tokens: int,
684
685
686
687
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
688
        use_cache: bool = True,
689
        **kwargs: Any,
690
    ) -> list[TokensTextLogprobs]:
691
692
693
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
694

695
696
697
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
698

699
        for inputs in all_inputs:
700
            output: "GenerateOutput" = self.model.generate(
701
                **self.wrap_device(inputs),
702
                use_cache=use_cache,
703
704
705
706
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
707
                **kwargs,
708
709
            )

710
711
712
713
714
715
            # Encoder-decoder models return decoder_hidden_states instead of
            # hidden_states
            hidden_states = (
                getattr(output, "hidden_states", None) or output.decoder_hidden_states
            )

716
717
718
            (
                seq_logprobs_lst,
                output_len,
719
            ) = self._hidden_states_to_logprobs(hidden_states, num_logprobs)
720
721
722
723
724
725
726

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

728
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
729
730
731
732
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
733

734
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
735
        return self.model.encode(prompts, *args, **kwargs)
736

737
738
    def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor:
        return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs)
739

740
741
742
743
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
744
        del self.model
745
        cleanup_dist_env_and_memory()
746

Woosuk Kwon's avatar
Woosuk Kwon committed
747

Cyrus Leung's avatar
Cyrus Leung committed
748
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
749
750
751
752
753
def hf_runner():
    return HfRunner


class VllmRunner:
754
755
    """
    The default value of some arguments have been modified from
756
    {class}`~vllm.LLM` as follows:
757

758
759
760
    - `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.
761
762
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
763
764
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
765
    - `enforce_eager`: Set to `False` to test CUDA graph.
766
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
767
768
769
770

    def __init__(
        self,
        model_name: str,
771
772
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
773
        tokenizer_name: str | None = None,
774
        tokenizer_mode: str = "auto",
775
        trust_remote_code: bool = True,
776
        seed: int = 0,
777
        max_model_len: int | None = 1024,
778
        dtype: str = "auto",
779
        disable_log_stats: bool = True,
780
        tensor_parallel_size: int = 1,
781
        block_size: int = 16 if not torch.xpu.is_available() else 64,
782
        enable_chunked_prefill: bool | None = False,
783
        swap_space: int = 4,
784
        enforce_eager: bool | None = False,
785
        # Set this to avoid hanging issue
786
        default_torch_num_threads: int | None = None,
787
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
788
    ) -> None:
789
790
791
792
793
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
794

795
        if not kwargs.get("compilation_config", None):
796
797
798
799
            # Note(@tdoublep): This is set to 4 because some tests (e.g., hybrid
            # model tests) may set max_num_seqs=4. If min cudagraph_capture_size is
            # set to larger than max_num_seqs, then it will lead to *no* graphs
            # being captured which can trigger edge cases that we don't handle yet.
800
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
801

802
803
804
805
806
807
808
809
            # Make sure we have atleast one cudagraph large enough for a single decode.
            if (speculative_config := kwargs.get("speculative_config")) and (
                num_speculative_tokens := speculative_config["num_speculative_tokens"]
            ):
                kwargs["compilation_config"]["cudagraph_capture_sizes"].append(
                    num_speculative_tokens + 1
                )

810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
        with init_ctx:
            self.llm = LLM(
                model=model_name,
                runner=runner,
                convert=convert,
                tokenizer=tokenizer_name,
                tokenizer_mode=tokenizer_mode,
                trust_remote_code=trust_remote_code,
                dtype=dtype,
                seed=seed,
                swap_space=swap_space,
                enforce_eager=enforce_eager,
                disable_log_stats=disable_log_stats,
                tensor_parallel_size=tensor_parallel_size,
                max_model_len=max_model_len,
                block_size=block_size,
                enable_chunked_prefill=enable_chunked_prefill,
                **kwargs,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
829

830
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
831
        self,
832
833
834
835
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
836
    ) -> list[dict[str, Any]]:
837
838
839
        if any(
            x is not None and len(x) != len(prompts) for x in [images, videos, audios]
        ):
840
            raise ValueError(
841
842
                "All non-None multimodal inputs must have the same length as prompts"
            )
843

844
        inputs = list[dict[str, Any]]()
845
        for i, prompt in enumerate(prompts):
846
847
848
849
850
851
852
853
854
            prompt_dict = dict[str, Any]()
            if isinstance(prompt, str):
                prompt_dict["prompt"] = prompt
            elif isinstance(prompt, list):
                prompt_dict["prompt_token_ids"] = prompt
            else:
                prompt_dict["prompt_embeds"] = prompt

            multi_modal_data = dict[str, Any]()
855
856
857
858
859
860
861
            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

862
863
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
864

865
            inputs.append(prompt_dict)
866
867
868
869
870

        return inputs

    def generate(
        self,
871
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
872
        sampling_params: SamplingParams,
873
874
875
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
876
        return_logprobs: bool = False,
877
        **kwargs: Any,
878
    ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
879
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
880

881
882
883
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
884

885
        outputs: list[tuple[list[list[int]], list[str]]] = []
886
        logprobs = []
Woosuk Kwon's avatar
Woosuk Kwon committed
887
888
889
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
890
891
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
892
            req_logprobs = []
893
894
            for sample in req_output.outputs:
                output_str = sample.text
895
                output_ids = list(sample.token_ids)
896
                req_sample_output_ids.append(prompt_ids + output_ids)
897
                req_sample_output_strs.append((prompt_str or "") + output_str)
898
899
                if sample.logprobs:
                    req_logprobs.extend(sample.logprobs)
900
            outputs.append((req_sample_output_ids, req_sample_output_strs))
901
902
            logprobs.append(req_logprobs)
        return outputs if not return_logprobs else (outputs, logprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
903

904
    @staticmethod
905
    def _final_steps_generate_w_logprobs(
906
        req_outputs: list[RequestOutput],
907
        include_prompt_token_ids: bool = False,
908
909
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
910
        for req_output in req_outputs:
911
            assert len(req_output.outputs) > 0
912
913
            for sample in req_output.outputs:
                output_str = sample.text
914
                output_ids = list(sample.token_ids)
915
                output_logprobs = sample.logprobs
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
            if include_prompt_token_ids:
                outputs.append(
                    (  # type: ignore[arg-type]
                        output_ids,
                        output_str,
                        output_logprobs,
                        req_output.prompt_token_ids,
                        req_output.prompt_logprobs,
                    )
                )
            else:
                outputs.append(
                    (
                        output_ids,
                        output_str,
                        output_logprobs,
                        req_output.prompt_logprobs,
                    )
                )

936
937
        return outputs

938
939
    def generate_w_logprobs(
        self,
940
        prompts: list[str],
941
        sampling_params: SamplingParams,
942
943
944
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
945
        include_prompt_token_ids: bool = False,
946
        **kwargs: Any,
947
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
948
949
950
951
952
953
954
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )

        toks_str_logsprobs_prompt_logprobs = self._final_steps_generate_w_logprobs(
955
            req_outputs, include_prompt_token_ids
956
        )
957
        # Omit prompt logprobs if not required by sampling params
958
959
960
961
962
        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
        )
963

Woosuk Kwon's avatar
Woosuk Kwon committed
964
965
    def generate_greedy(
        self,
966
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
967
        max_tokens: int,
968
969
970
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
971
        **kwargs: Any,
972
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
973
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
974
975
976
977
978
979
980
981
982
        outputs = self.generate(
            prompts,
            greedy_params,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
983

984
985
    def generate_greedy_logprobs(
        self,
986
        prompts: list[str],
987
        max_tokens: int,
988
989
990
991
992
993
994
        num_logprobs: int | None,
        num_prompt_logprobs: int | None = None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
        stop_token_ids: list[int] | None = None,
        stop: list[str] | None = None,
995
        **kwargs: Any,
996
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
997
998
999
1000
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
1001
            prompt_logprobs=num_prompt_logprobs,
1002
            stop_token_ids=stop_token_ids,
1003
1004
            stop=stop,
        )
1005

1006
1007
1008
1009
1010
1011
1012
1013
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
1014

1015
1016
1017
1018
1019
1020
1021
    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
        """
1022
1023
1024
        outputs = self.generate_greedy_logprobs(
            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
        )
1025
1026
1027
1028

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
1029
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
            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

1043
    def generate_beam_search(
1044
        self,
1045
        prompts: list[str],
1046
1047
        beam_width: int,
        max_tokens: int,
1048
1049
1050
1051
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
1052
    ) -> list[tuple[list[list[int]], list[str]]]:
1053
1054
1055
1056
1057
1058
1059
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

        outputs = self.llm.beam_search(
            inputs,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens),
            concurrency_limit=concurrency_limit,
        )
1060
1061
1062
1063
1064
1065
1066
        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

1067
    def classify(self, prompts: list[str]) -> list[list[float]]:
1068
        req_outputs = self.llm.classify(prompts)
1069
1070
        return [req_output.outputs.probs for req_output in req_outputs]

1071
1072
1073
    def embed(
        self,
        prompts: list[str],
1074
1075
1076
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1077
1078
1079
1080
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1081

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

1085
1086
1087
1088
1089
1090
    def token_embed(self, prompts: list[str]) -> list[list[float]]:
        req_outputs = self.llm.encode(prompts, pooling_task="token_embed")
        return [req_output.outputs.data for req_output in req_outputs]

    def token_classify(self, prompts: list[str]) -> list[list[float]]:
        req_outputs = self.llm.encode(prompts, pooling_task="token_classify")
1091
1092
        return [req_output.outputs.data for req_output in req_outputs]

1093
1094
1095
1096
    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]

1097
1098
    def score(
        self,
1099
1100
        text_1: list[str] | str,
        text_2: list[str] | str,
1101
1102
        *args,
        **kwargs,
1103
    ) -> list[float]:
1104
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1105
        return [req_output.outputs.score for req_output in req_outputs]
1106

1107
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1108
        return self.llm.apply_model(func)
1109

1110
1111
1112
    def get_llm(self) -> LLM:
        return self.llm

1113
1114
1115
1116
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1117
        del self.llm
1118
        cleanup_dist_env_and_memory()
1119

Woosuk Kwon's avatar
Woosuk Kwon committed
1120

1121
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1122
1123
def vllm_runner():
    return VllmRunner
1124
1125


1126
1127
1128
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1129

1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
    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
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
@pytest.fixture()
def caplog_mp_fork():
    """
    This fixture enables capturing logs from a forked MP subprocess.
    It should be used in conjunction with caplog_vllm.

    By default, subprocess logs do not go through the parent process.
    We instead create a queue listener in the parent process which
    forwards logs to the logger's other handlers, and add a QueueHandler
    to the root logger. Forked subprocesses will inherit the root logger
    and pass their messages to the queue, which the listener will forward
    to the root logger, which can be captured by caplog.

    Note that this workaround only works for fork; with spawn, the subprocess
    reinitializes logging and does not automatically inherit the queue.
    We'd have to manually pass the queue to the subprocess at the spawn point.
    See caplog_mp_spawn below.
    """

    @contextlib.contextmanager
    def ctx():
        import logging.handlers
        import multiprocessing as mp

        logger_queue: mp.Queue[logging.LogRecord] = mp.Queue()
        logger = logging.getLogger()
        handlers = logger.handlers

        # The listener works on a background thread, not inherited by the child.
        queue_listener = logging.handlers.QueueListener(logger_queue, *handlers)
        queue_listener.start()

        # Add queue handler after creating the listener to avoid cycle
        logger.addHandler(logging.handlers.QueueHandler(logger_queue))
        yield
        queue_listener.stop()

    return ctx


class LogHolder:
    def __init__(self):
        self.text = None


@pytest.fixture()
def caplog_mp_spawn(tmp_path, monkeypatch):
    """
    This fixture enables capturing logs from a forked MP subprocess.
    It does not require caplog_vllm (but it only contains logs from the child).

    By default, subprocess logs do not go through the parent process.
    We instead add a FileHandler to the config so the spawned child process
    writes its logs to a temp file.
    In the parent, we read the file and return the contents.

    Note: this method could be extended to fork by either reconfiguring logging
    in the parent or using a SocketHandler:
    https://docs.python.org/3/howto/logging-cookbook.html#sending-and-receiving-logging-events-across-a-network # noqa: E501
    """

    @contextlib.contextmanager
    def ctx(level: int | str):
        from vllm.logger import DEFAULT_LOGGING_CONFIG

        config_path = tmp_path / "vllm_logging_config.json"
        log_path = tmp_path / "vllm.log"
        log_holder = LogHolder()

        config = deepcopy(DEFAULT_LOGGING_CONFIG)
        if envs.VLLM_LOGGING_CONFIG_PATH:
            path = pathlib.Path(envs.VLLM_LOGGING_CONFIG_PATH)
            assert path.exists()
            config = json.loads(path.read_text())

        config["loggers"]["vllm"]["handlers"] += ["vllm_file"]
        config["handlers"]["vllm_file"] = {
            "class": "logging.FileHandler",
            "formatter": "vllm",
            "level": level,
            "filename": log_path.as_posix(),
        }
1225
        config["loggers"]["vllm"]["level"] = level
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238

        config_path.write_text(json.dumps(config))

        with monkeypatch.context() as monkeypatch_ctx:
            monkeypatch_ctx.setenv("VLLM_LOGGING_CONFIG_PATH", config_path.as_posix())
            monkeypatch_ctx.setenv("VLLM_CONFIGURE_LOGGING", "1")
            yield log_holder

        log_holder.text = log_path.read_text()

    return ctx


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

1244
    from vllm.platforms import current_platform
1245

1246
    return current_platform.device_count()
1247
1248
1249


temp_dir = tempfile.gettempdir()
1250
1251
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1252
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1253
1254
1255
1256


@pytest.fixture
def dummy_opt_path():
1257
1258
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1259
1260
1261
1262
1263
        snapshot_download(
            repo_id="facebook/opt-125m",
            local_dir=_dummy_opt_path,
            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
        )
1264
        assert os.path.exists(json_path)
1265
        with open(json_path) as f:
1266
1267
1268
1269
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1270
1271
1272
1273
1274
1275
1276
    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):
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
        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",
                "*.safetensors",
            ],
        )
1289
        assert os.path.exists(json_path)
1290
        with open(json_path) as f:
1291
1292
1293
1294
1295
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1296
1297
1298
1299
1300
1301


@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):
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1314
        assert os.path.exists(json_path)
1315
        with open(json_path) as f:
1316
1317
1318
1319
1320
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1321
1322
1323
1324
1325


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1326
1327
1328
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1329
1330
1331
1332
1333
1334
1335
1336
1337


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:
1338
            item.add_marker(skip_optional)
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350


@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")
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397


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
1398
1399
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409

    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(
1410
1411
1412
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
        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:
1440
1441
1442
1443
1444
        image = fetch_image(self.url_for(name))
        # Unwrap MediaWithBytes if present
        if isinstance(image, MediaWithBytes):
            image = image.media
        return image
1445
1446
1447
1448
1449


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
1450
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
    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]
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482


@pytest.fixture
def disable_deepgemm_ue8m0(monkeypatch):
    from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

    with monkeypatch.context() as monkeypatch_ctx:
        monkeypatch_ctx.setenv("VLLM_USE_DEEP_GEMM_E8M0", "0")
        is_deep_gemm_e8m0_used.cache_clear()
        yield
        # Clear cache so the next time it is used it is processed with the
        # default VLLM_USE_DEEP_GEMM_E8M0  setting.
        is_deep_gemm_e8m0_used.cache_clear()
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511


@pytest.fixture(autouse=True)
def clean_gpu_memory_between_tests():
    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
        yield
        return

    # Wait for GPU memory to be cleared before starting the test
    import gc

    from tests.utils import wait_for_gpu_memory_to_clear

    num_gpus = torch.cuda.device_count()
    if num_gpus > 0:
        try:
            wait_for_gpu_memory_to_clear(
                devices=list(range(num_gpus)),
                threshold_ratio=0.1,
            )
        except ValueError as e:
            logger.info("Failed to clean GPU memory: %s", e)

    yield

    # Clean up GPU memory after the test
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522


@pytest.fixture
def use_fresh_inductor_cache():
    """
    Use a fresh inductor cache for the test.
    This is useful to ensure that the test is not affected by the
    previous test calls.
    """
    with fresh_cache():
        yield