conftest.py 52.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
chenzk's avatar
chenzk committed
3
4
5
6
7
8
import os

# Keep v1 CUDA-graph tests valid: fork default for ``LLM(kvprune_compression=None)``
# is controlled by ``VLLM_KVPRUNE_COMPRESSION_DEFAULT`` (see ``env_override``).
os.environ.setdefault("VLLM_KVPRUNE_COMPRESSION_DEFAULT", "0")
os.environ.setdefault("VLLM_KVPRUNE_RELEASE_V1_KV", "0")
9
10
11
import contextlib
import pathlib
from copy import deepcopy
12
13
14

from tblib import pickling_support

15
16
17
# Import fixture
from tests.v1.entrypoints.conftest import sample_json_schema  # noqa

18
19
# ruff: noqa

20
21
22
23
24
# 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()

25
import http.server
26
import json
27
import math
28
import mimetypes
29
import os
30
import socket
31
import tempfile
32
33
import threading
from collections.abc import Generator
34
from contextlib import nullcontext
35
from enum import Enum
36
from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING, Optional
Woosuk Kwon's avatar
Woosuk Kwon committed
37

38
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
39
40
import pytest
import torch
41
import torch.nn as nn
42
import torch.nn.functional as F
43
from huggingface_hub import snapshot_download
44
from PIL import Image
45
46
47
48
49
50
51
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BatchEncoding,
    BatchFeature,
)
52
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
53

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

80
81
82
from torch._inductor.utils import fresh_cache


83
84
85
86
87
if TYPE_CHECKING:
    from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
    from transformers.generation.utils import GenerateOutput


88
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
89

90
91
92
_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")]
93
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
94

Cyrus Leung's avatar
Cyrus Leung committed
95
_M = TypeVar("_M")
96

97
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
98
99

PromptImageInput = _PromptMultiModalInput[Image.Image]
100
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
101
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
102

103

104
def _read_prompts(filename: str) -> list[str]:
105
    with open(filename) as f:
106
107
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
108
109


110
class ImageAssetPrompts(TypedDict):
111
112
    stop_sign: str
    cherry_blossom: str
113
114


115
class ImageTestAssets(list[ImageAsset]):
116
    def __init__(self) -> None:
117
118
119
120
121
122
        super().__init__(
            [
                ImageAsset("stop_sign"),
                ImageAsset("cherry_blossom"),
            ]
        )
123

124
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
125
126
127
128
129
130
        """
        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.
        """
131
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
132
133


134
135
class VideoAssetPrompts(TypedDict):
    baby_reading: str
136
137


138
class VideoTestAssets(list[VideoAsset]):
139
    def __init__(self) -> None:
140
141
142
143
144
        super().__init__(
            [
                VideoAsset("baby_reading"),
            ]
        )
145

146
147
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
148
149


150
class AudioAssetPrompts(TypedDict):
151
152
153
154
    mary_had_lamb: str
    winning_call: str


155
class AudioTestAssets(list[AudioAsset]):
156
    def __init__(self) -> None:
157
158
159
160
161
162
        super().__init__(
            [
                AudioAsset("mary_had_lamb"),
                AudioAsset("winning_call"),
            ]
        )
163

164
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
165
166
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

167

168
IMAGE_ASSETS = ImageTestAssets()
169
"""Singleton instance of {class}`ImageTestAssets`."""
170
VIDEO_ASSETS = VideoTestAssets()
171
"""Singleton instance of {class}`VideoTestAssets`."""
172
AUDIO_ASSETS = AudioTestAssets()
173
"""Singleton instance of {class}`AudioTestAssets`."""
174
175


176
177
178
179
180
181
182
@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


183
184
@pytest.fixture
def dist_init():
185
186
    from tests.utils import ensure_current_vllm_config

187
    temp_file = tempfile.mkstemp()[1]
188
189
190
191
192
193
194
195
196
197
198

    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
199
    cleanup_dist_env_and_memory()
200
201


202
203
204
205
206
207
208
209
210
211
212
@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

    with set_current_vllm_config(VllmConfig()):
        yield


213
@pytest.fixture()
214
def should_do_global_cleanup_after_test(request) -> bool:
215
216
217
218
    """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.
    """
219

220
    return not request.node.get_closest_marker("skip_global_cleanup")
221
222


223
@pytest.fixture(autouse=True)
224
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
225
    yield
226
    if should_do_global_cleanup_after_test:
227
        cleanup_dist_env_and_memory()
228
229


230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
@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()


251
252
253
254
255
256
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
257
@pytest.fixture
258
def example_prompts() -> list[str]:
259
    return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
260
261


262
263
264
265
266
267
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


268
269
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
270

271
272
273
274
275
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


276
@pytest.fixture
277
def example_long_prompts() -> list[str]:
278
    return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
Woosuk Kwon's avatar
Woosuk Kwon committed
279
280


281
@pytest.fixture(scope="session")
282
def image_assets() -> ImageTestAssets:
283
284
285
    return IMAGE_ASSETS


286
@pytest.fixture(scope="session")
287
def video_assets() -> VideoTestAssets:
288
289
290
    return VIDEO_ASSETS


291
@pytest.fixture(scope="session")
292
def audio_assets() -> AudioTestAssets:
293
294
295
    return AUDIO_ASSETS


296
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
297
_R = TypeVar("_R")
298

Woosuk Kwon's avatar
Woosuk Kwon committed
299
300

class HfRunner:
301
    def get_default_device(self):
302
        from vllm.platforms import current_platform
303

304
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
305

306
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
307
        if x is None or isinstance(x, (bool,)):
308
309
            return x

310
        if device is None:
311
            device = self.device
312

313
314
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
315

316
317
318
319
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
320

Woosuk Kwon's avatar
Woosuk Kwon committed
321
322
323
    def __init__(
        self,
        model_name: str,
324
        dtype: str = "auto",
325
        *,
326
        model_kwargs: dict[str, Any] | None = None,
327
        trust_remote_code: bool = True,
328
        is_sentence_transformer: bool = False,
329
        is_cross_encoder: bool = False,
330
        skip_tokenizer_init: bool = False,
331
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
332
        # Set this to avoid hanging issue
333
        default_torch_num_threads: int | None = None,
334
    ) -> None:
335
336
337
338
339
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357

        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",
        *,
358
        model_kwargs: dict[str, Any] | None = None,
359
360
361
362
363
        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
364
    ) -> None:
365
        model_name = maybe_model_redirect(model_name)
366
        self.model_name = model_name
367

368
369
        self.config = AutoConfig.from_pretrained(
            model_name,
370
            trust_remote_code=trust_remote_code,
371
372
        )
        self.device = self.get_default_device()
373
        self.dtype = dtype = _get_and_verify_dtype(
374
375
376
377
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
378
            config_format="hf",
379
        )
380
381

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

384
        if is_sentence_transformer:
385
386
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
387
388
389
390
391

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
392
                trust_remote_code=trust_remote_code,
393
            )
394
395
396
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
397
398
399
400
401

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
402
                trust_remote_code=trust_remote_code,
403
            )
404
        else:
405
406
407
408
409
410
411
            model = cast(
                nn.Module,
                auto_cls.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                    **model_kwargs,
                ),
412
413
            )

414
            # in case some unquantized custom models are not in same dtype
415
416
417
            if getattr(model, "quantization_method", None) is None and any(
                p.dtype != self.dtype for p in model.parameters()
            ):
418
419
                model = model.to(dtype=self.dtype)

420
421
422
423
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
424
                model = model.to(device=self.device)
425
426

            self.model = model
427

428
        if not skip_tokenizer_init:
429
430
431
432
433
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                )
434
            )
435

436
        # don't put this import at the top level
437
        # it will call torch.accelerator.device_count()
438
        from transformers import AutoProcessor
439

440
441
        self.processor = AutoProcessor.from_pretrained(
            model_name,
442
            trust_remote_code=trust_remote_code,
443
        )
444
445
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
446

447
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
448
        self,
449
450
451
452
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
453
        tokenization_kwargs: dict[str, Any] | None = None,
454
    ) -> list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]]:
455
        if images is not None:
456
            assert len(prompts) == len(images)
457

458
459
460
461
462
463
        if videos is not None:
            assert len(prompts) == len(videos)

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

464
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
465
        for i, prompt in enumerate(prompts):
466
            if isinstance(prompt, str):
467
468
469
470
471
472
473
474
475
476
477
478
                # 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",
                    }
                )
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
                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)
512
513
514

        return all_inputs

515
516
517
518
519
520
521
522
523
    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

524
    def classify(self, prompts: list[str]) -> list[list[float]]:
525
526
        # output is final logits
        all_inputs = self.get_inputs(prompts)
527
        outputs: list[list[float]] = []
528
529
        problem_type = getattr(self.config, "problem_type", "")

530
531
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
532
533
534

            assert isinstance(output.logits, torch.Tensor)

535
536
537
538
539
            if problem_type == "regression":
                logits = output.logits[0].tolist()
            elif problem_type == "multi_label_classification":
                logits = output.logits.sigmoid()[0].tolist()
            else:
540
                logits = softmax(output.logits)[0].tolist()
541
542
543
544
            outputs.append(logits)

        return outputs

545
546
    def generate(
        self,
547
548
549
550
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
551
        **kwargs: Any,
552
    ) -> list[tuple[list[list[int]], list[str]]]:
553
554
555
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
556

557
        outputs: list[tuple[list[list[int]], list[str]]] = []
558
        for inputs in all_inputs:
559
            output_ids: torch.Tensor = self.model.generate(
560
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
561
562
563
                use_cache=True,
                **kwargs,
            )
564
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
565
566
567
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
568
            )
569
            outputs.append((output_ids.cpu().tolist(), output_str))
Woosuk Kwon's avatar
Woosuk Kwon committed
570
571
572
573
        return outputs

    def generate_greedy(
        self,
574
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
575
        max_tokens: int,
576
577
578
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
579
        **kwargs: Any,
580
    ) -> list[tuple[list[int], str]]:
581
582
583
584
585
586
587
588
589
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
590

591
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
592
593
594

    def generate_beam_search(
        self,
595
        prompts: list[str],
596
597
        beam_width: int,
        max_tokens: int,
598
599
600
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
601
    ) -> list[tuple[list[list[int]], list[str]]]:
602
603
604
605
606
607
608
609
610
611
        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,
        )
612

613
614
615
616
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
617
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
618
619
620
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
621

622
623
    def generate_greedy_logprobs(
        self,
624
        prompts: list[str],
625
        max_tokens: int,
626
627
628
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
629
        **kwargs: Any,
630
    ) -> list[list[torch.Tensor]]:
631
632
633
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
634

635
        all_logprobs: list[list[torch.Tensor]] = []
636
        for inputs in all_inputs:
637
            output: "GenerateOutput" = self.model.generate(
638
                **self.wrap_device(inputs),
639
640
641
642
643
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
644
                **kwargs,
645
            )
646
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
647
648
649
            all_logprobs.append(seq_logprobs)
        return all_logprobs

650
    def _hidden_states_to_seq_logprobs(
651
        self,
652
653
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
654
655
        output_embeddings = self.model.get_output_embeddings()

656
        seq_logprobs: list[torch.Tensor] = []
657
658
659
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
660
661
662
663
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
664
                output_embeddings.weight.t(),
665
            )
666
667
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
668
669
670
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

671
672
673
674
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
675
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
676
        num_logprobs: int | None,
677
    ) -> tuple[list[dict[int, float]], int]:
678
679
680
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

681
        # convert to dict
682
        seq_logprobs_lst: list[dict[int, float]] = []
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
        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,
        )

700
701
    def generate_greedy_logprobs_limit(
        self,
702
        prompts: list[str],
703
        max_tokens: int,
704
705
706
707
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
708
        use_cache: bool = True,
709
        **kwargs: Any,
710
    ) -> list[TokensTextLogprobs]:
711
712
713
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
714

715
716
717
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
718

719
        for inputs in all_inputs:
720
            output: "GenerateOutput" = self.model.generate(
721
                **self.wrap_device(inputs),
722
                use_cache=use_cache,
723
724
725
726
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
727
                **kwargs,
728
729
            )

730
731
732
733
734
735
            # Encoder-decoder models return decoder_hidden_states instead of
            # hidden_states
            hidden_states = (
                getattr(output, "hidden_states", None) or output.decoder_hidden_states
            )

736
737
738
            (
                seq_logprobs_lst,
                output_len,
739
            ) = self._hidden_states_to_logprobs(hidden_states, num_logprobs)
740
741
742
743
744
745
746

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

748
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
749
750
751
752
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
753

754
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
755
        return self.model.encode(prompts, *args, **kwargs)
756

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

760
761
762
763
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
764
        del self.model
765
        cleanup_dist_env_and_memory()
766

Woosuk Kwon's avatar
Woosuk Kwon committed
767

Cyrus Leung's avatar
Cyrus Leung committed
768
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
769
770
771
772
773
def hf_runner():
    return HfRunner


class VllmRunner:
774
775
    """
    The default value of some arguments have been modified from
776
    {class}`~vllm.LLM` as follows:
777

778
779
780
    - `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.
781
782
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
783
784
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
785
    - `enforce_eager`: Set to `False` to test CUDA graph.
786
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
787
788
789
790

    def __init__(
        self,
        model_name: str,
791
792
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
793
        tokenizer_name: str | None = None,
794
        tokenizer_mode: str = "auto",
795
        trust_remote_code: bool = True,
796
        seed: int = 0,
797
        max_model_len: int | None = 1024,
798
        dtype: str = "auto",
799
        disable_log_stats: bool = True,
800
        tensor_parallel_size: int = 1,
801
        block_size: int = 16 if not torch.xpu.is_available() else 64,
802
803
        enable_chunked_prefill: bool | None = False,
        enforce_eager: bool | None = False,
804
        # Set this to avoid hanging issue
805
        default_torch_num_threads: int | None = None,
806
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
807
    ) -> None:
808
809
810
811
812
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
813

814
        if not kwargs.get("compilation_config", None):
815
816
817
818
            # 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.
819
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
820

821
822
823
824
825
826
827
828
            # 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
                )

829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
        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
847

848
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
849
        self,
850
851
852
853
        prompts: list[str]
        | list[torch.Tensor]
        | list[list[int]]
        | list[dict[str, Any]],
854
855
856
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
857
    ) -> list[dict[str, Any]]:
858
859
860
        if any(
            x is not None and len(x) != len(prompts) for x in [images, videos, audios]
        ):
861
            raise ValueError(
862
863
                "All non-None multimodal inputs must have the same length as prompts"
            )
864

865
        inputs = list[dict[str, Any]]()
866
        for i, prompt in enumerate(prompts):
867
868
869
870
871
            # 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())
872
            else:
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
                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
888

889
890
                if multi_modal_data:
                    prompt_dict["multi_modal_data"] = multi_modal_data
891

892
                inputs.append(prompt_dict)
893
894
895
896
897

        return inputs

    def generate(
        self,
898
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
899
        sampling_params: SamplingParams,
900
901
902
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
903
        return_logprobs: bool = False,
904
        **kwargs: Any,
905
    ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
906
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
907

908
909
910
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
911

912
        outputs: list[tuple[list[list[int]], list[str]]] = []
913
        logprobs = []
Woosuk Kwon's avatar
Woosuk Kwon committed
914
915
916
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
917
918
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
919
            req_logprobs = []
920
921
            for sample in req_output.outputs:
                output_str = sample.text
922
                output_ids = list(sample.token_ids)
923
                req_sample_output_ids.append(prompt_ids + output_ids)
924
                req_sample_output_strs.append((prompt_str or "") + output_str)
925
926
                if sample.logprobs:
                    req_logprobs.extend(sample.logprobs)
927
            outputs.append((req_sample_output_ids, req_sample_output_strs))
928
929
            logprobs.append(req_logprobs)
        return outputs if not return_logprobs else (outputs, logprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
930

931
    @staticmethod
932
    def _final_steps_generate_w_logprobs(
933
        req_outputs: list[RequestOutput],
934
        include_prompt_token_ids: bool = False,
935
936
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
937
        for req_output in req_outputs:
938
            assert len(req_output.outputs) > 0
939
940
            for sample in req_output.outputs:
                output_str = sample.text
941
                output_ids = list(sample.token_ids)
942
                output_logprobs = sample.logprobs
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
            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,
                    )
                )

963
964
        return outputs

965
966
    def generate_w_logprobs(
        self,
967
        prompts: list[str],
968
        sampling_params: SamplingParams,
969
970
971
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
972
        include_prompt_token_ids: bool = False,
973
        **kwargs: Any,
974
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
975
976
977
978
979
980
981
        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(
982
            req_outputs, include_prompt_token_ids
983
        )
984
        # Omit prompt logprobs if not required by sampling params
985
986
987
988
989
        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
        )
990

Woosuk Kwon's avatar
Woosuk Kwon committed
991
992
    def generate_greedy(
        self,
993
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
994
        max_tokens: int,
995
996
997
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
998
        **kwargs: Any,
999
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
1000
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
1001
1002
1003
1004
1005
1006
1007
1008
1009
        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]
1010

1011
1012
    def generate_greedy_logprobs(
        self,
1013
        prompts: list[str],
1014
        max_tokens: int,
1015
1016
1017
1018
1019
1020
1021
        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,
1022
        **kwargs: Any,
1023
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
1024
1025
1026
1027
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
1028
            prompt_logprobs=num_prompt_logprobs,
1029
            stop_token_ids=stop_token_ids,
1030
1031
            stop=stop,
        )
1032

1033
1034
1035
1036
1037
1038
1039
1040
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
1041

1042
1043
1044
    def generate_prompt_perplexity(
        self, prompts: list[str], mask: Optional[list[str]] = None
    ) -> list[float]:
1045
1046
1047
1048
1049
1050
        """
        Return the perplexity score associated with generating the prompts

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

1055
1056
1057
1058
1059
1060
        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))]
        )

1061
        perplexities = []
1062
        for output, mask_prefix_len in zip(outputs, mask_prefix_lens):
1063
            output = cast(TokensTextLogprobsPromptLogprobs, output)
1064
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
1065
            assert token_datas[0] is None
1066

1067
            token_log_probs = []
1068
            for token_data in token_datas[mask_prefix_len + 1 :]:
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
                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

1079
    def generate_beam_search(
1080
        self,
1081
        prompts: list[str],
1082
1083
        beam_width: int,
        max_tokens: int,
1084
1085
1086
1087
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
1088
    ) -> list[tuple[list[list[int]], list[str]]]:
1089
1090
1091
1092
1093
1094
1095
        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,
        )
1096
1097
1098
1099
1100
1101
1102
        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

1103
    def classify(self, prompts: list[str]) -> list[list[float]]:
1104
        req_outputs = self.llm.classify(prompts)
1105
1106
        return [req_output.outputs.probs for req_output in req_outputs]

1107
1108
1109
    def embed(
        self,
        prompts: list[str],
1110
1111
1112
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1113
1114
1115
1116
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1117

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

1121
1122
1123
1124
1125
1126
    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")
1127
1128
        return [req_output.outputs.data for req_output in req_outputs]

1129
1130
1131
1132
    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]

1133
1134
    def score(
        self,
1135
1136
        text_1: list[str] | str,
        text_2: list[str] | str,
1137
1138
        *args,
        **kwargs,
1139
    ) -> list[float]:
1140
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1141
        return [req_output.outputs.score for req_output in req_outputs]
1142

1143
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1144
        return self.llm.apply_model(func)
1145

1146
1147
1148
    def get_llm(self) -> LLM:
        return self.llm

1149
1150
1151
    def collective_rpc(self, *args, **kwargs):
        return self.llm.collective_rpc(*args, **kwargs)

1152
1153
1154
1155
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1156
1157
1158
1159
1160
1161
1162
1163
1164
        # 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
1165
        del self.llm
1166
        cleanup_dist_env_and_memory()
1167

Woosuk Kwon's avatar
Woosuk Kwon committed
1168

1169
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1170
1171
def vllm_runner():
    return VllmRunner
1172
1173


1174
1175
1176
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1177

1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
    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
1189
1190


1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
@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(),
        }
1273
        config["loggers"]["vllm"]["level"] = level
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286

        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


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

1292
    from vllm.platforms import current_platform
1293

1294
    return current_platform.device_count()
1295
1296
1297


temp_dir = tempfile.gettempdir()
1298
1299
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1300
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1301
1302
1303
1304


@pytest.fixture
def dummy_opt_path():
1305
1306
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1307
1308
1309
1310
1311
        snapshot_download(
            repo_id="facebook/opt-125m",
            local_dir=_dummy_opt_path,
            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
        )
1312
        assert os.path.exists(json_path)
1313
        with open(json_path) as f:
1314
1315
1316
1317
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1318
1319
1320
1321
1322
1323
1324
    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):
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
        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",
            ],
        )
1337
        assert os.path.exists(json_path)
1338
        with open(json_path) as f:
1339
1340
1341
1342
1343
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1344
1345
1346
1347
1348
1349


@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):
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1362
        assert os.path.exists(json_path)
1363
        with open(json_path) as f:
1364
1365
1366
1367
1368
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1369
1370
1371
1372
1373


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1374
1375
1376
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1377
1378
1379
1380
1381
1382
1383
1384
1385


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:
1386
            item.add_marker(skip_optional)
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398


@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")
1399
1400
1401
1402
1403
1404
1405
1406
1407
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
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445


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
1446
1447
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457

    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(
1458
1459
1460
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
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
        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:
1488
1489
1490
1491
1492
        image = fetch_image(self.url_for(name))
        # Unwrap MediaWithBytes if present
        if isinstance(image, MediaWithBytes):
            image = image.media
        return image
1493
1494
1495
1496
1497


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
1498
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
    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]
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530


@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()
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543


@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

1544
    num_gpus = torch.accelerator.device_count()
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
    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():
1558
        torch.accelerator.empty_cache()
1559
        gc.collect()
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570


@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
1571
1572


1573
1574
1575
1576
1577
1578
1579
1580
@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


1581
1582
1583
1584
@pytest.fixture(scope="function")
def enable_pickle(monkeypatch):
    """`LLM.apply_model` requires pickling a function."""
    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")