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

70
71
72
from torch._inductor.utils import fresh_cache


73
74
75
76
77
if TYPE_CHECKING:
    from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
    from transformers.generation.utils import GenerateOutput


78
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
79

80
81
82
_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")]
83
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
84

Cyrus Leung's avatar
Cyrus Leung committed
85
_M = TypeVar("_M")
86

87
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
88
89

PromptImageInput = _PromptMultiModalInput[Image.Image]
90
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
91
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
92

93

94
def _read_prompts(filename: str) -> list[str]:
95
    with open(filename) as f:
96
97
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
98
99


100
class ImageAssetPrompts(TypedDict):
101
102
    stop_sign: str
    cherry_blossom: str
103
104


105
class ImageTestAssets(list[ImageAsset]):
106
    def __init__(self) -> None:
107
108
109
110
111
112
        super().__init__(
            [
                ImageAsset("stop_sign"),
                ImageAsset("cherry_blossom"),
            ]
        )
113

114
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
115
116
117
118
119
120
        """
        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.
        """
121
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
122
123


124
125
class VideoAssetPrompts(TypedDict):
    baby_reading: str
126
127


128
class VideoTestAssets(list[VideoAsset]):
129
    def __init__(self) -> None:
130
131
132
133
134
        super().__init__(
            [
                VideoAsset("baby_reading"),
            ]
        )
135

136
137
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
138
139


140
class AudioAssetPrompts(TypedDict):
141
142
143
144
    mary_had_lamb: str
    winning_call: str


145
class AudioTestAssets(list[AudioAsset]):
146
    def __init__(self) -> None:
147
148
149
150
151
152
        super().__init__(
            [
                AudioAsset("mary_had_lamb"),
                AudioAsset("winning_call"),
            ]
        )
153

154
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
155
156
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

157

158
IMAGE_ASSETS = ImageTestAssets()
159
"""Singleton instance of {class}`ImageTestAssets`."""
160
VIDEO_ASSETS = VideoTestAssets()
161
"""Singleton instance of {class}`VideoTestAssets`."""
162
AUDIO_ASSETS = AudioTestAssets()
163
"""Singleton instance of {class}`AudioTestAssets`."""
164
165


166
167
168
169
170
171
172
@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


173
174
175
176
177
178
179
180
181
182
183
184
@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
185
    cleanup_dist_env_and_memory()
186
187


188
@pytest.fixture()
189
def should_do_global_cleanup_after_test(request) -> bool:
190
191
192
193
    """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.
    """
194

195
    return not request.node.get_closest_marker("skip_global_cleanup")
196
197


198
@pytest.fixture(autouse=True)
199
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
200
    yield
201
    if should_do_global_cleanup_after_test:
202
        cleanup_dist_env_and_memory()
203
204


205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
@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()


226
227
228
229
230
231
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
232
@pytest.fixture
233
def example_prompts() -> list[str]:
234
    return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
235
236


237
238
239
240
241
242
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


243
244
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
245

246
247
248
249
250
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


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


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


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


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


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

Woosuk Kwon's avatar
Woosuk Kwon committed
274
275

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

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

281
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
282
        if x is None or isinstance(x, (bool,)):
283
284
            return x

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

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

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

        return x.to(device)
295

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

        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",
        *,
333
        model_kwargs: dict[str, Any] | None = None,
334
335
336
337
338
        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
339
    ) -> None:
340
        model_name = maybe_model_redirect(model_name)
341
        self.model_name = model_name
342

343
344
        self.config = AutoConfig.from_pretrained(
            model_name,
345
            trust_remote_code=trust_remote_code,
346
347
        )
        self.device = self.get_default_device()
348
        self.dtype = dtype = _get_and_verify_dtype(
349
350
351
352
353
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
354
355

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

358
        if is_sentence_transformer:
359
360
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
361
362
363
364
365

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

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
376
                trust_remote_code=trust_remote_code,
377
            )
378
        else:
379
380
381
382
383
384
385
            model = cast(
                nn.Module,
                auto_cls.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                    **model_kwargs,
                ),
386
387
            )

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

394
395
396
397
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
398
                model = model.to(device=self.device)
399
400

            self.model = model
401

402
        if not skip_tokenizer_init:
403
404
405
406
407
408
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
                    dtype=dtype,
                    trust_remote_code=trust_remote_code,
                )
409
            )
410

411
412
413
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
414

415
416
        self.processor = AutoProcessor.from_pretrained(
            model_name,
417
            dtype=dtype,
418
            trust_remote_code=trust_remote_code,
419
        )
420
421
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
422

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

434
435
436
437
438
439
        if videos is not None:
            assert len(prompts) == len(videos)

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

440
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
441
        for i, prompt in enumerate(prompts):
442
            if isinstance(prompt, str):
443
444
445
446
447
448
449
450
451
452
453
454
                # 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",
                    }
                )
455
456
457
458
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
                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)
488
489
490

        return all_inputs

491
492
493
494
495
496
497
498
499
    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

500
    def classify(self, prompts: list[str]) -> list[list[float]]:
501
502
        # output is final logits
        all_inputs = self.get_inputs(prompts)
503
        outputs: list[list[float]] = []
504
505
        problem_type = getattr(self.config, "problem_type", "")

506
507
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
508
509
510

            assert isinstance(output.logits, torch.Tensor)

511
512
513
514
515
516
            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()
517
518
519
520
            outputs.append(logits)

        return outputs

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

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

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

567
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
568
569
570

    def generate_beam_search(
        self,
571
        prompts: list[str],
572
573
        beam_width: int,
        max_tokens: int,
574
575
576
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
577
    ) -> list[tuple[list[list[int]], list[str]]]:
578
579
580
581
582
583
584
585
586
587
        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,
        )
588

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

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

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

626
    def _hidden_states_to_seq_logprobs(
627
        self,
628
629
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
630
631
        output_embeddings = self.model.get_output_embeddings()

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

647
648
649
650
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
651
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
652
        num_logprobs: int | None,
653
    ) -> tuple[list[dict[int, float]], int]:
654
655
656
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

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

676
677
    def generate_greedy_logprobs_limit(
        self,
678
        prompts: list[str],
679
        max_tokens: int,
680
681
682
683
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
684
        **kwargs: Any,
685
    ) -> list[TokensTextLogprobs]:
686
687
688
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
689

690
691
692
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
693

694
        for inputs in all_inputs:
695
            output: "GenerateOutput" = self.model.generate(
696
                **self.wrap_device(inputs),
697
698
699
700
701
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
702
                **kwargs,
703
704
            )

705
706
707
708
709
710
            # Encoder-decoder models return decoder_hidden_states instead of
            # hidden_states
            hidden_states = (
                getattr(output, "hidden_states", None) or output.decoder_hidden_states
            )

711
712
713
            (
                seq_logprobs_lst,
                output_len,
714
            ) = self._hidden_states_to_logprobs(hidden_states, num_logprobs)
715
716
717
718
719
720
721

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

723
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
724
725
726
727
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
728

729
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
730
        return self.model.encode(prompts, *args, **kwargs)
731

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

735
736
737
738
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
739
        del self.model
740
        cleanup_dist_env_and_memory()
741

Woosuk Kwon's avatar
Woosuk Kwon committed
742

Cyrus Leung's avatar
Cyrus Leung committed
743
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
744
745
746
747
748
def hf_runner():
    return HfRunner


class VllmRunner:
749
750
    """
    The default value of some arguments have been modified from
751
    {class}`~vllm.LLM` as follows:
752

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

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

790
        if not kwargs.get("compilation_config", None):
791
792
793
794
            # 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.
795
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
796

797
798
799
800
801
802
803
804
            # 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
                )

805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
        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
824

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

839
        inputs = list[dict[str, Any]]()
840
        for i, prompt in enumerate(prompts):
841
842
843
844
845
846
847
848
849
            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]()
850
851
852
853
854
855
856
            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

857
858
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
859

860
            inputs.append(prompt_dict)
861
862
863
864
865

        return inputs

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

876
877
878
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
879

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

899
    @staticmethod
900
    def _final_steps_generate_w_logprobs(
901
        req_outputs: list[RequestOutput],
902
        include_prompt_token_ids: bool = False,
903
904
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
905
        for req_output in req_outputs:
906
            assert len(req_output.outputs) > 0
907
908
            for sample in req_output.outputs:
                output_str = sample.text
909
                output_ids = list(sample.token_ids)
910
                output_logprobs = sample.logprobs
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
            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,
                    )
                )

931
932
        return outputs

933
934
    def generate_w_logprobs(
        self,
935
        prompts: list[str],
936
        sampling_params: SamplingParams,
937
938
939
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
940
        include_prompt_token_ids: bool = False,
941
        **kwargs: Any,
942
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
943
944
945
946
947
948
949
        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(
950
            req_outputs, include_prompt_token_ids
951
        )
952
        # Omit prompt logprobs if not required by sampling params
953
954
955
956
957
        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
        )
958

Woosuk Kwon's avatar
Woosuk Kwon committed
959
960
    def generate_greedy(
        self,
961
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
962
        max_tokens: int,
963
964
965
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
966
        **kwargs: Any,
967
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
968
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
969
970
971
972
973
974
975
976
977
        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]
978

979
980
    def generate_greedy_logprobs(
        self,
981
        prompts: list[str],
982
        max_tokens: int,
983
984
985
986
987
988
989
        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,
990
        **kwargs: Any,
991
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
992
993
994
995
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
996
            prompt_logprobs=num_prompt_logprobs,
997
            stop_token_ids=stop_token_ids,
998
999
            stop=stop,
        )
1000

1001
1002
1003
1004
1005
1006
1007
1008
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
1009

1010
1011
1012
1013
1014
1015
1016
    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
        """
1017
1018
1019
        outputs = self.generate_greedy_logprobs(
            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
        )
1020
1021
1022
1023

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
1024
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
            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

1038
    def generate_beam_search(
1039
        self,
1040
        prompts: list[str],
1041
1042
        beam_width: int,
        max_tokens: int,
1043
1044
1045
1046
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
1047
    ) -> list[tuple[list[list[int]], list[str]]]:
1048
1049
1050
1051
1052
1053
1054
        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,
        )
1055
1056
1057
1058
1059
1060
1061
        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

1062
    def classify(self, prompts: list[str]) -> list[list[float]]:
1063
        req_outputs = self.llm.classify(prompts)
1064
1065
        return [req_output.outputs.probs for req_output in req_outputs]

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

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

1080
1081
1082
1083
1084
1085
    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")
1086
1087
        return [req_output.outputs.data for req_output in req_outputs]

1088
1089
1090
1091
    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]

1092
1093
    def score(
        self,
1094
1095
        text_1: list[str] | str,
        text_2: list[str] | str,
1096
1097
        *args,
        **kwargs,
1098
    ) -> list[float]:
1099
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1100
        return [req_output.outputs.score for req_output in req_outputs]
1101

1102
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1103
        return self.llm.apply_model(func)
1104

1105
1106
1107
    def get_llm(self) -> LLM:
        return self.llm

1108
1109
1110
1111
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1112
        del self.llm
1113
        cleanup_dist_env_and_memory()
1114

Woosuk Kwon's avatar
Woosuk Kwon committed
1115

1116
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1117
1118
def vllm_runner():
    return VllmRunner
1119
1120


1121
1122
1123
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1124

1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
    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
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
@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(),
        }
1220
        config["loggers"]["vllm"]["level"] = level
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233

        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


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

1239
    from vllm.platforms import current_platform
1240

1241
    return current_platform.device_count()
1242
1243
1244


temp_dir = tempfile.gettempdir()
1245
1246
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1247
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1248
1249
1250
1251


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


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


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


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:
1333
            item.add_marker(skip_optional)
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345


@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")
1346
1347
1348
1349
1350
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


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
1393
1394
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404

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


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


@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()
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506


@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()
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517


@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