conftest.py 55 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
# ruff: noqa

11
12
13
14
15
# 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()

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

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

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

71
72
73
from torch._inductor.utils import fresh_cache


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


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

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

@pytest.fixture
def sample_json_schema():
    return {
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"},
            "skills": {
                "type": "array",
                "items": {
                    "type": "string",
                },
            },
            "grade": {
                "type": "string",
                "pattern": "^[A-D]$",
            },
            "email": {
                "type": "string",
                "pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$",
            },
            "work_history": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "company": {"type": "string"},
                        "duration": {
                            "type": "number",
                            "minimum": 0.0,
                            "maximum": 100.0,
                        },
                        "position": {"type": "string"},
                    },
                    "required": ["company", "duration", "position"],
                    "additionalProperties": False,
                },
                "minItems": 0,
                "maxItems": 3,
            },
        },
        "required": ["name", "age", "skills", "grade", "email", "work_history"],
        "additionalProperties": False,
        "minProperties": 1,
        "maxProperties": 10,
    }


130
131
132
_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")]
133
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
134

Cyrus Leung's avatar
Cyrus Leung committed
135
_M = TypeVar("_M")
136

137
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
138
139

PromptImageInput = _PromptMultiModalInput[Image.Image]
140
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
141
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
142

143

144
def _read_prompts(filename: str) -> list[str]:
145
    with open(filename) as f:
146
147
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
148
149


150
class ImageAssetPrompts(TypedDict):
151
152
    stop_sign: str
    cherry_blossom: str
153
154


155
class ImageTestAssets(list[ImageAsset]):
156
    def __init__(self) -> None:
157
158
159
160
161
162
        super().__init__(
            [
                ImageAsset("stop_sign"),
                ImageAsset("cherry_blossom"),
            ]
        )
163

164
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
165
166
167
168
169
170
        """
        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.
        """
171
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
172
173


174
175
class VideoAssetPrompts(TypedDict):
    baby_reading: str
176
177


178
class VideoTestAssets(list[VideoAsset]):
179
    def __init__(self) -> None:
180
181
182
183
184
        super().__init__(
            [
                VideoAsset("baby_reading"),
            ]
        )
185

186
187
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
188
189


190
class AudioAssetPrompts(TypedDict):
191
192
193
194
    mary_had_lamb: str
    winning_call: str


195
class AudioTestAssets(list[AudioAsset]):
196
    def __init__(self) -> None:
197
198
199
200
201
202
        super().__init__(
            [
                AudioAsset("mary_had_lamb"),
                AudioAsset("winning_call"),
            ]
        )
203

204
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
205
206
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

207

208
IMAGE_ASSETS = ImageTestAssets()
209
"""Singleton instance of {class}`ImageTestAssets`."""
210
VIDEO_ASSETS = VideoTestAssets()
211
"""Singleton instance of {class}`VideoTestAssets`."""
212
AUDIO_ASSETS = AudioTestAssets()
213
"""Singleton instance of {class}`AudioTestAssets`."""
214
215


216
217
218
219
220
221
222
@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


223
224
@pytest.fixture
def dist_init():
225
226
    from tests.utils import ensure_current_vllm_config

227
    temp_file = tempfile.mkstemp()[1]
228
229
230
231
232
233
234
235
236
237
238

    with ensure_current_vllm_config():
        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
239
    cleanup_dist_env_and_memory()
240
241


242
243
244
245
246
247
248
@pytest.fixture
def default_vllm_config():
    """Set a default VllmConfig for tests that directly test CustomOps or pathways
    that use get_current_vllm_config() outside of a full engine context.
    """
    from vllm.config import VllmConfig, set_current_vllm_config

249
250
251
    config = VllmConfig()
    with set_current_vllm_config(config):
        yield config
252
253


254
@pytest.fixture()
255
def should_do_global_cleanup_after_test(request) -> bool:
256
257
258
259
    """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.
    """
260

261
    return not request.node.get_closest_marker("skip_global_cleanup")
262
263


264
@pytest.fixture(autouse=True)
265
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
266
    yield
267
    if should_do_global_cleanup_after_test:
268
        cleanup_dist_env_and_memory()
269
270


271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
@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()


292
293
294
295
296
297
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
298
@pytest.fixture
299
def example_prompts() -> list[str]:
300
    return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
301
302


303
304
305
306
307
308
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


309
310
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
311

312
313
314
315
316
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


317
@pytest.fixture
318
def example_long_prompts() -> list[str]:
319
    return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
Woosuk Kwon's avatar
Woosuk Kwon committed
320
321


322
@pytest.fixture(scope="session")
323
def image_assets() -> ImageTestAssets:
324
325
326
    return IMAGE_ASSETS


327
@pytest.fixture(scope="session")
328
def video_assets() -> VideoTestAssets:
329
330
331
    return VIDEO_ASSETS


332
@pytest.fixture(scope="session")
333
def audio_assets() -> AudioTestAssets:
334
335
336
    return AUDIO_ASSETS


337
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
338
_R = TypeVar("_R")
339

Woosuk Kwon's avatar
Woosuk Kwon committed
340
341

class HfRunner:
342
    def get_default_device(self):
343
        from vllm.platforms import current_platform
344

345
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
346

347
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
348
        if x is None or isinstance(x, (bool,)):
349
350
            return x

351
        if device is None:
352
            device = self.device
353

354
355
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
356

357
358
359
360
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
361

Woosuk Kwon's avatar
Woosuk Kwon committed
362
363
364
    def __init__(
        self,
        model_name: str,
365
        dtype: str = "auto",
366
        *,
367
        revision: str | None = None,
368
        model_kwargs: dict[str, Any] | None = None,
369
        trust_remote_code: bool = True,
370
        is_sentence_transformer: bool = False,
371
        is_cross_encoder: bool = False,
372
        skip_tokenizer_init: bool = False,
373
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
374
        # Set this to avoid hanging issue
375
        default_torch_num_threads: int | None = None,
376
    ) -> None:
377
378
379
380
381
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
382
383
384
385
386

        with init_ctx:
            self._init(
                model_name=model_name,
                dtype=dtype,
387
                revision=revision,
388
389
390
391
392
393
394
395
396
397
398
399
400
                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",
        *,
401
        revision: str | None = None,
402
        model_kwargs: dict[str, Any] | None = None,
403
404
405
406
407
        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
408
    ) -> None:
409
        model_name = maybe_model_redirect(model_name)
410
        self.model_name = model_name
411

412
413
        self.config = AutoConfig.from_pretrained(
            model_name,
414
            trust_remote_code=trust_remote_code,
415
        )
416
417
418
419
420
421
422
423
424
        # HF runner should use the HF config so that it's consistent with the HF model
        if self.config.__module__.startswith("vllm.transformers_utils.configs"):
            from transformers.models.auto.configuration_auto import CONFIG_MAPPING

            del CONFIG_MAPPING._extra_content[self.config.model_type]
            self.config = AutoConfig.from_pretrained(
                model_name,
                trust_remote_code=trust_remote_code,
            )
425
        self.device = self.get_default_device()
426
        self.dtype = dtype = _get_and_verify_dtype(
427
428
429
430
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
431
            config_format="hf",
432
        )
433
434

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

437
        if is_sentence_transformer:
438
439
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
440
441
442

            self.model = SentenceTransformer(
                model_name,
443
                revision=revision,
444
445
                device=self.device,
                model_kwargs=model_kwargs,
446
                trust_remote_code=trust_remote_code,
447
            )
448
449
450
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
451
452
453

            self.model = CrossEncoder(
                model_name,
454
                revision=revision,
455
456
                device=self.device,
                automodel_args=model_kwargs,
457
                trust_remote_code=trust_remote_code,
458
            )
459
        else:
460
461
462
463
            model = cast(
                nn.Module,
                auto_cls.from_pretrained(
                    model_name,
464
                    revision=revision,
465
466
467
                    trust_remote_code=trust_remote_code,
                    **model_kwargs,
                ),
468
469
            )

470
            # in case some unquantized custom models are not in same dtype
471
472
473
            if getattr(model, "quantization_method", None) is None and any(
                p.dtype != self.dtype for p in model.parameters()
            ):
474
475
                model = model.to(dtype=self.dtype)

476
477
478
479
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
480
                model = model.to(device=self.device)
481
482

            self.model = model
483

484
        if not skip_tokenizer_init:
485
486
487
488
489
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                )
490
            )
491

492
        # don't put this import at the top level
493
        # it will call torch.accelerator.device_count()
494
        from transformers import AutoProcessor
495

496
497
        self.processor = AutoProcessor.from_pretrained(
            model_name,
498
            trust_remote_code=trust_remote_code,
499
        )
500
501
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
502

503
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
504
        self,
505
506
507
508
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
509
        tokenization_kwargs: dict[str, Any] | None = None,
510
    ) -> list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]]:
511
        if images is not None:
512
            assert len(prompts) == len(images)
513

514
515
516
517
518
519
        if videos is not None:
            assert len(prompts) == len(videos)

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

520
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
521
        for i, prompt in enumerate(prompts):
522
            if isinstance(prompt, str):
523
524
525
526
527
528
529
530
531
532
533
534
                # 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",
                    }
                )
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
                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)
568
569
570

        return all_inputs

571
572
573
574
575
576
577
578
579
    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

580
    def classify(self, prompts: list[str]) -> list[list[float]]:
581
582
        # output is final logits
        all_inputs = self.get_inputs(prompts)
583
        outputs: list[list[float]] = []
584
585
        problem_type = getattr(self.config, "problem_type", "")

586
587
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
588
589
590

            assert isinstance(output.logits, torch.Tensor)

591
592
593
594
595
            if problem_type == "regression":
                logits = output.logits[0].tolist()
            elif problem_type == "multi_label_classification":
                logits = output.logits.sigmoid()[0].tolist()
            else:
596
                logits = softmax(output.logits)[0].tolist()
597
598
599
600
            outputs.append(logits)

        return outputs

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

613
        outputs: list[tuple[list[list[int]], list[str]]] = []
614
        for inputs in all_inputs:
615
            output_ids: torch.Tensor = self.model.generate(
616
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
617
618
619
                use_cache=True,
                **kwargs,
            )
620
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
621
622
623
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
624
            )
625
            outputs.append((output_ids.cpu().tolist(), output_str))
Woosuk Kwon's avatar
Woosuk Kwon committed
626
627
628
629
        return outputs

    def generate_greedy(
        self,
630
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
631
        max_tokens: int,
632
633
634
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
635
        **kwargs: Any,
636
    ) -> list[tuple[list[int], str]]:
637
638
639
640
641
642
643
644
645
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
646

647
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
648
649
650

    def generate_beam_search(
        self,
651
        prompts: list[str],
652
653
        beam_width: int,
        max_tokens: int,
654
655
656
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
657
    ) -> list[tuple[list[list[int]], list[str]]]:
658
659
660
661
662
663
664
665
666
667
        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,
        )
668

669
670
671
672
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
673
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
674
675
676
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
677

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

691
        all_logprobs: list[list[torch.Tensor]] = []
692
        for inputs in all_inputs:
693
            output: "GenerateOutput" = self.model.generate(
694
                **self.wrap_device(inputs),
695
696
697
698
699
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
700
                **kwargs,
701
            )
702
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
703
704
705
            all_logprobs.append(seq_logprobs)
        return all_logprobs

706
    def _hidden_states_to_seq_logprobs(
707
        self,
708
709
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
710
711
        output_embeddings = self.model.get_output_embeddings()

712
        seq_logprobs: list[torch.Tensor] = []
713
714
715
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
716
717
718
719
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
720
                output_embeddings.weight.t(),
721
            )
722
723
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
724
725
726
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

727
728
729
730
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
731
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
732
        num_logprobs: int | None,
733
    ) -> tuple[list[dict[int, float]], int]:
734
735
736
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

737
        # convert to dict
738
        seq_logprobs_lst: list[dict[int, float]] = []
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
        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,
        )

756
757
    def generate_greedy_logprobs_limit(
        self,
758
        prompts: list[str],
759
        max_tokens: int,
760
761
762
763
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
764
        use_cache: bool = True,
765
        **kwargs: Any,
766
    ) -> list[TokensTextLogprobs]:
767
768
769
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
770

771
772
773
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
774

775
        for inputs in all_inputs:
776
            output: "GenerateOutput" = self.model.generate(
777
                **self.wrap_device(inputs),
778
                use_cache=use_cache,
779
780
781
782
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
783
                **kwargs,
784
785
            )

786
787
788
789
790
791
            # Encoder-decoder models return decoder_hidden_states instead of
            # hidden_states
            hidden_states = (
                getattr(output, "hidden_states", None) or output.decoder_hidden_states
            )

792
793
794
            (
                seq_logprobs_lst,
                output_len,
795
            ) = self._hidden_states_to_logprobs(hidden_states, num_logprobs)
796
797
798
799
800
801
802

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

804
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
805
806
807
808
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
809

810
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
811
        return self.model.encode(prompts, *args, **kwargs)
812

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

816
817
818
819
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
820
        del self.model
821
        cleanup_dist_env_and_memory()
822

Woosuk Kwon's avatar
Woosuk Kwon committed
823

Cyrus Leung's avatar
Cyrus Leung committed
824
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
825
826
827
828
829
def hf_runner():
    return HfRunner


class VllmRunner:
830
831
    """
    The default value of some arguments have been modified from
832
    {class}`~vllm.LLM` as follows:
833

834
835
836
    - `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.
837
838
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
839
840
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
841
    - `enforce_eager`: Set to `False` to test CUDA graph.
842
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
843
844
845
846

    def __init__(
        self,
        model_name: str,
847
848
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
849
        tokenizer_name: str | None = None,
850
        tokenizer_mode: str = "auto",
851
        trust_remote_code: bool = True,
852
        seed: int = 0,
853
        max_model_len: int | None = 1024,
854
        dtype: str = "auto",
855
        disable_log_stats: bool = True,
856
        tensor_parallel_size: int = 1,
857
        block_size: int = 16 if not torch.xpu.is_available() else 64,
858
859
        enable_chunked_prefill: bool | None = False,
        enforce_eager: bool | None = False,
860
        # Set this to avoid hanging issue
861
        default_torch_num_threads: int | None = None,
862
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
863
    ) -> None:
864
865
866
867
868
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
869

870
        if not kwargs.get("compilation_config", None):
871
872
873
874
            # 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.
875
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
876

877
878
879
880
881
882
883
884
            # 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
                )

885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
        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,
                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
903

904
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
905
        self,
906
907
908
909
        prompts: list[str]
        | list[torch.Tensor]
        | list[list[int]]
        | list[dict[str, Any]],
910
911
912
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
913
    ) -> list[dict[str, Any]]:
914
915
916
        if any(
            x is not None and len(x) != len(prompts) for x in [images, videos, audios]
        ):
917
            raise ValueError(
918
919
                "All non-None multimodal inputs must have the same length as prompts"
            )
920

921
        inputs = list[dict[str, Any]]()
922
        for i, prompt in enumerate(prompts):
923
924
925
926
927
            # If we're passing an encoder/decoder prompt, we assume it
            # already contains the multimodal data in the prompt
            if isinstance(prompt, dict):
                assert images is None and audios is None and videos is None
                inputs.append(prompt.copy())
928
            else:
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
                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]()
                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
944

945
946
                if multi_modal_data:
                    prompt_dict["multi_modal_data"] = multi_modal_data
947

948
                inputs.append(prompt_dict)
949
950
951
952
953

        return inputs

    def generate(
        self,
954
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
955
        sampling_params: SamplingParams,
956
957
958
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
959
        return_logprobs: bool = False,
960
        **kwargs: Any,
961
    ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
962
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
963

964
965
966
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
967

968
        outputs: list[tuple[list[list[int]], list[str]]] = []
969
        logprobs = []
Woosuk Kwon's avatar
Woosuk Kwon committed
970
971
972
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
973
974
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
975
            req_logprobs = []
976
977
            for sample in req_output.outputs:
                output_str = sample.text
978
                output_ids = list(sample.token_ids)
979
                req_sample_output_ids.append(prompt_ids + output_ids)
980
                req_sample_output_strs.append((prompt_str or "") + output_str)
981
982
                if sample.logprobs:
                    req_logprobs.extend(sample.logprobs)
983
            outputs.append((req_sample_output_ids, req_sample_output_strs))
984
985
            logprobs.append(req_logprobs)
        return outputs if not return_logprobs else (outputs, logprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
986

987
    @staticmethod
988
    def _final_steps_generate_w_logprobs(
989
        req_outputs: list[RequestOutput],
990
        include_prompt_token_ids: bool = False,
991
992
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
993
        for req_output in req_outputs:
994
            assert len(req_output.outputs) > 0
995
996
            for sample in req_output.outputs:
                output_str = sample.text
997
                output_ids = list(sample.token_ids)
998
                output_logprobs = sample.logprobs
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
            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,
                    )
                )

1019
1020
        return outputs

1021
1022
    def generate_w_logprobs(
        self,
1023
        prompts: list[str],
1024
        sampling_params: SamplingParams,
1025
1026
1027
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
1028
        include_prompt_token_ids: bool = False,
1029
        **kwargs: Any,
1030
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
1031
1032
1033
1034
1035
1036
1037
        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(
1038
            req_outputs, include_prompt_token_ids
1039
        )
1040
        # Omit prompt logprobs if not required by sampling params
1041
1042
1043
1044
1045
        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
        )
1046

Woosuk Kwon's avatar
Woosuk Kwon committed
1047
1048
    def generate_greedy(
        self,
1049
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
1050
        max_tokens: int,
1051
1052
1053
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1054
        **kwargs: Any,
1055
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
1056
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
1057
1058
1059
1060
1061
1062
1063
1064
1065
        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]
1066

1067
1068
    def generate_greedy_logprobs(
        self,
1069
        prompts: list[str],
1070
        max_tokens: int,
1071
1072
1073
1074
1075
1076
1077
        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,
1078
        **kwargs: Any,
1079
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
1080
1081
1082
1083
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
1084
            prompt_logprobs=num_prompt_logprobs,
1085
            stop_token_ids=stop_token_ids,
1086
1087
            stop=stop,
        )
1088

1089
1090
1091
1092
1093
1094
1095
1096
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
1097

1098
1099
1100
    def generate_prompt_perplexity(
        self, prompts: list[str], mask: Optional[list[str]] = None
    ) -> list[float]:
1101
1102
1103
1104
1105
1106
        """
        Return the perplexity score associated with generating the prompts

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

1111
1112
1113
1114
1115
1116
        mask_prefix_lens = (
            [len(self.llm.get_tokenizer()(prefix)["input_ids"]) for prefix in mask]
            if mask is not None
            else [0 for _ in range(len(prompts))]
        )

1117
        perplexities = []
1118
        for output, mask_prefix_len in zip(outputs, mask_prefix_lens):
1119
            output = cast(TokensTextLogprobsPromptLogprobs, output)
1120
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
1121
            assert token_datas[0] is None
1122

1123
            token_log_probs = []
1124
            for token_data in token_datas[mask_prefix_len + 1 :]:
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
                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

1135
    def generate_beam_search(
1136
        self,
1137
        prompts: list[str],
1138
1139
        beam_width: int,
        max_tokens: int,
1140
1141
1142
1143
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
1144
    ) -> list[tuple[list[list[int]], list[str]]]:
1145
1146
1147
1148
1149
1150
1151
        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,
        )
1152
1153
1154
1155
1156
1157
1158
        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

1159
    def classify(self, prompts: list[str]) -> list[list[float]]:
1160
        req_outputs = self.llm.classify(prompts)
1161
1162
        return [req_output.outputs.probs for req_output in req_outputs]

1163
1164
1165
    def embed(
        self,
        prompts: list[str],
1166
1167
1168
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1169
1170
1171
1172
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1173

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

1177
1178
1179
1180
1181
1182
    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")
1183
1184
        return [req_output.outputs.data for req_output in req_outputs]

1185
1186
1187
1188
    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]

1189
1190
    def score(
        self,
1191
1192
        text_1: list[str] | str,
        text_2: list[str] | str,
1193
1194
        *args,
        **kwargs,
1195
    ) -> list[float]:
1196
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1197
        return [req_output.outputs.score for req_output in req_outputs]
1198

1199
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1200
        return self.llm.apply_model(func)
1201

1202
1203
1204
    def get_llm(self) -> LLM:
        return self.llm

1205
1206
1207
    def collective_rpc(self, *args, **kwargs):
        return self.llm.collective_rpc(*args, **kwargs)

1208
1209
1210
1211
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1212
1213
1214
1215
1216
1217
1218
1219
1220
        # Explicitly shutdown the engine core to release GPU resources
        # This is needed because when executing consecutive tests, the GC
        # might not be fast enough in shutting down the llm engine. This can lead to OOMs
        # because when the next test starts some GPU memory is still in use.
        try:
            self.llm.llm_engine.engine_core.shutdown()
        except Exception:
            # Ignore shutdown errors as cleanup will still proceed
            pass
1221
        del self.llm
1222
        cleanup_dist_env_and_memory()
1223

Woosuk Kwon's avatar
Woosuk Kwon committed
1224

1225
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1226
1227
def vllm_runner():
    return VllmRunner
1228
1229


1230
1231
1232
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1233

1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
    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
1245
1246


1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
@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(),
        }
1329
        config["loggers"]["vllm"]["level"] = level
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342

        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


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

1348
    from vllm.platforms import current_platform
1349

1350
    return current_platform.device_count()
1351
1352
1353


temp_dir = tempfile.gettempdir()
1354
1355
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1356
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1357
1358
1359
1360


@pytest.fixture
def dummy_opt_path():
1361
1362
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1363
1364
1365
1366
1367
        snapshot_download(
            repo_id="facebook/opt-125m",
            local_dir=_dummy_opt_path,
            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
        )
1368
        assert os.path.exists(json_path)
1369
        with open(json_path) as f:
1370
1371
1372
1373
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1374
1375
1376
1377
1378
1379
1380
    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):
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
        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",
            ],
        )
1393
        assert os.path.exists(json_path)
1394
        with open(json_path) as f:
1395
1396
1397
1398
1399
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1400
1401
1402
1403
1404
1405


@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):
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1418
        assert os.path.exists(json_path)
1419
        with open(json_path) as f:
1420
1421
1422
1423
1424
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1425
1426
1427
1428
1429


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1430
1431
1432
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1433
1434
1435
1436
1437
1438
1439
1440
1441


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:
1442
            item.add_marker(skip_optional)
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454


@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")
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501


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
1502
1503
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513

    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(
1514
1515
1516
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
        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:
1544
1545
1546
1547
1548
        image = fetch_image(self.url_for(name))
        # Unwrap MediaWithBytes if present
        if isinstance(image, MediaWithBytes):
            image = image.media
        return image
1549
1550
1551
1552
1553


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
1554
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
    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]
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586


@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()
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599


@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

1600
    num_gpus = torch.accelerator.device_count()
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
    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():
1614
        torch.accelerator.empty_cache()
1615
        gc.collect()
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626


@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
1627
1628


1629
1630
1631
1632
1633
1634
1635
1636
@pytest.fixture
def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache):
    """Temporary VLLM_CACHE_ROOT combined with a fresh inductor cache."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        monkeypatch.setenv("VLLM_CACHE_ROOT", tmp_dir)
        yield tmp_dir


1637
1638
1639
1640
@pytest.fixture(scope="function")
def enable_pickle(monkeypatch):
    """`LLM.apply_model` requires pickling a function."""
    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663


@pytest.fixture(scope="function")
def disable_log_dedup(monkeypatch):
    """
    Disable log deduplication such that warning_once and info_once always print.
    """

    # Patch logger._print_warning_once to remove the lru_cache decorator
    from vllm import logger

    original_print_warning_once = logger._print_warning_once
    original_print_info_once = logger._print_info_once
    original_print_debug_once = logger._print_debug_once

    logger._print_warning_once = original_print_warning_once.__wrapped__
    logger._print_info_once = original_print_info_once.__wrapped__
    logger._print_debug_once = original_print_debug_once.__wrapped__

    yield
    logger._print_warning_once = original_print_warning_once
    logger._print_info_once = original_print_info_once
    logger._print_debug_once = original_print_debug_once