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

from tblib import pickling_support

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

12
13
# ruff: noqa

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

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

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

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

70
71
72
73
74
if TYPE_CHECKING:
    from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
    from transformers.generation.utils import GenerateOutput


75
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
76

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

Cyrus Leung's avatar
Cyrus Leung committed
82
_M = TypeVar("_M")
83

84
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
85
86

PromptImageInput = _PromptMultiModalInput[Image.Image]
87
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
88
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
89

90

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


97
class ImageAssetPrompts(TypedDict):
98
99
    stop_sign: str
    cherry_blossom: str
100
101


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

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


121
122
class VideoAssetPrompts(TypedDict):
    baby_reading: str
123
124


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

133
134
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
135
136


137
class AudioAssetPrompts(TypedDict):
138
139
140
141
    mary_had_lamb: str
    winning_call: str


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

151
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
152
153
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

154

155
IMAGE_ASSETS = ImageTestAssets()
156
"""Singleton instance of {class}`ImageTestAssets`."""
157
VIDEO_ASSETS = VideoTestAssets()
158
"""Singleton instance of {class}`VideoTestAssets`."""
159
AUDIO_ASSETS = AudioTestAssets()
160
"""Singleton instance of {class}`AudioTestAssets`."""
161
162


163
164
165
166
167
168
169
@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


170
171
172
173
174
175
176
177
178
179
180
181
@pytest.fixture
def dist_init():
    temp_file = tempfile.mkstemp()[1]
    init_distributed_environment(
        world_size=1,
        rank=0,
        distributed_init_method=f"file://{temp_file}",
        local_rank=0,
        backend="nccl",
    )
    initialize_model_parallel(1, 1)
    yield
182
    cleanup_dist_env_and_memory()
183
184


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

192
    return not request.node.get_closest_marker("skip_global_cleanup")
193
194


195
@pytest.fixture(autouse=True)
196
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
197
    yield
198
    if should_do_global_cleanup_after_test:
199
        cleanup_dist_env_and_memory()
200
201


202
203
204
205
206
207
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
208
@pytest.fixture
209
def example_prompts() -> list[str]:
210
    return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)]
211
212


213
214
215
216
217
218
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


219
220
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
221

222
223
224
225
226
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


227
@pytest.fixture
228
def example_long_prompts() -> list[str]:
229
    return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
Woosuk Kwon's avatar
Woosuk Kwon committed
230
231


232
@pytest.fixture(scope="session")
233
def image_assets() -> ImageTestAssets:
234
235
236
    return IMAGE_ASSETS


237
@pytest.fixture(scope="session")
238
def video_assets() -> VideoTestAssets:
239
240
241
    return VIDEO_ASSETS


242
@pytest.fixture(scope="session")
243
def audio_assets() -> AudioTestAssets:
244
245
246
    return AUDIO_ASSETS


247
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
248
_R = TypeVar("_R")
249

Woosuk Kwon's avatar
Woosuk Kwon committed
250
251

class HfRunner:
252
    def get_default_device(self):
253
        from vllm.platforms import current_platform
254

255
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
256

257
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
258
        if x is None or isinstance(x, (bool,)):
259
260
            return x

261
        if device is None:
262
            device = self.device
263

264
265
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
266

267
268
269
270
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
271

Woosuk Kwon's avatar
Woosuk Kwon committed
272
273
274
    def __init__(
        self,
        model_name: str,
275
        dtype: str = "auto",
276
        *,
277
        model_kwargs: dict[str, Any] | None = None,
278
        trust_remote_code: bool = True,
279
        is_sentence_transformer: bool = False,
280
        is_cross_encoder: bool = False,
281
        skip_tokenizer_init: bool = False,
282
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
283
        # Set this to avoid hanging issue
284
        default_torch_num_threads: int | None = None,
285
    ) -> None:
286
287
288
289
290
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308

        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",
        *,
309
        model_kwargs: dict[str, Any] | None = None,
310
311
312
313
314
        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
315
    ) -> None:
316
        model_name = maybe_model_redirect(model_name)
317
        self.model_name = model_name
318

319
320
        self.config = AutoConfig.from_pretrained(
            model_name,
321
            trust_remote_code=trust_remote_code,
322
323
        )
        self.device = self.get_default_device()
324
        self.dtype = dtype = _get_and_verify_dtype(
325
326
327
328
329
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
330
331

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

334
        if is_sentence_transformer:
335
336
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
337
338
339
340
341

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
342
                trust_remote_code=trust_remote_code,
343
            )
344
345
346
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
347
348
349
350
351

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
352
                trust_remote_code=trust_remote_code,
353
            )
354
        else:
355
356
357
358
359
360
361
            model = cast(
                nn.Module,
                auto_cls.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code,
                    **model_kwargs,
                ),
362
363
            )

364
            # in case some unquantized custom models are not in same dtype
365
366
367
            if getattr(model, "quantization_method", None) is None and any(
                p.dtype != self.dtype for p in model.parameters()
            ):
368
369
                model = model.to(dtype=self.dtype)

370
371
372
373
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
374
                model = model.to(device=self.device)
375
376

            self.model = model
377

378
        if not skip_tokenizer_init:
379
380
381
382
383
384
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
                    dtype=dtype,
                    trust_remote_code=trust_remote_code,
                )
385
            )
386

387
388
389
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
390

391
392
        self.processor = AutoProcessor.from_pretrained(
            model_name,
393
            dtype=dtype,
394
            trust_remote_code=trust_remote_code,
395
        )
396
397
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
398

399
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
400
        self,
401
402
403
404
405
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
    ) -> list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]]:
406
        if images is not None:
407
            assert len(prompts) == len(images)
408

409
410
411
412
413
414
        if videos is not None:
            assert len(prompts) == len(videos)

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

415
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
416
        for i, prompt in enumerate(prompts):
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
            if isinstance(prompt, str):
                processor_kwargs: dict[str, Any] = {
                    "text": prompt,
                    "return_tensors": "pt",
                }
                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)
455
456
457

        return all_inputs

458
459
460
461
462
463
464
465
466
    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

467
    def classify(self, prompts: list[str]) -> list[list[float]]:
468
469
        # output is final logits
        all_inputs = self.get_inputs(prompts)
470
        outputs: list[list[float]] = []
471
472
        problem_type = getattr(self.config, "problem_type", "")

473
474
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
475
476
477

            assert isinstance(output.logits, torch.Tensor)

478
479
480
481
482
483
            if problem_type == "regression":
                logits = output.logits[0].tolist()
            elif problem_type == "multi_label_classification":
                logits = output.logits.sigmoid()[0].tolist()
            else:
                logits = output.logits.softmax(dim=-1)[0].tolist()
484
485
486
487
            outputs.append(logits)

        return outputs

488
489
    def generate(
        self,
490
491
492
493
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
494
        **kwargs: Any,
495
    ) -> list[tuple[list[list[int]], list[str]]]:
496
497
498
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
499

500
        outputs: list[tuple[list[list[int]], list[str]]] = []
501
        for inputs in all_inputs:
502
            output_ids: torch.Tensor = self.model.generate(
503
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
504
505
506
                use_cache=True,
                **kwargs,
            )
507
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
508
509
510
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
511
            )
512
            outputs.append((output_ids.cpu().tolist(), output_str))
Woosuk Kwon's avatar
Woosuk Kwon committed
513
514
515
516
        return outputs

    def generate_greedy(
        self,
517
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
518
        max_tokens: int,
519
520
521
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
522
        **kwargs: Any,
523
    ) -> list[tuple[list[int], str]]:
524
525
526
527
528
529
530
531
532
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
533

534
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
535
536
537

    def generate_beam_search(
        self,
538
        prompts: list[str],
539
540
        beam_width: int,
        max_tokens: int,
541
542
543
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
544
    ) -> list[tuple[list[list[int]], list[str]]]:
545
546
547
548
549
550
551
552
553
554
        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,
        )
555

556
557
558
559
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
560
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
561
562
563
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
564

565
566
    def generate_greedy_logprobs(
        self,
567
        prompts: list[str],
568
        max_tokens: int,
569
570
571
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
572
        **kwargs: Any,
573
    ) -> list[list[torch.Tensor]]:
574
575
576
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
577

578
        all_logprobs: list[list[torch.Tensor]] = []
579
        for inputs in all_inputs:
580
            output: "GenerateOutput" = self.model.generate(
581
                **self.wrap_device(inputs),
582
583
584
585
586
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
587
                **kwargs,
588
            )
589
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
590
591
592
            all_logprobs.append(seq_logprobs)
        return all_logprobs

593
    def _hidden_states_to_seq_logprobs(
594
        self,
595
596
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
597
598
        output_embeddings = self.model.get_output_embeddings()

599
        seq_logprobs: list[torch.Tensor] = []
600
601
602
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
603
604
605
606
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
607
                output_embeddings.weight.t(),
608
            )
609
610
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
611
612
613
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

614
615
616
617
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
618
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
619
        num_logprobs: int | None,
620
    ) -> tuple[list[dict[int, float]], int]:
621
622
623
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

624
        # convert to dict
625
        seq_logprobs_lst: list[dict[int, float]] = []
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
        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,
        )

643
644
    def generate_greedy_logprobs_limit(
        self,
645
        prompts: list[str],
646
        max_tokens: int,
647
648
649
650
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
651
        **kwargs: Any,
652
    ) -> list[TokensTextLogprobs]:
653
654
655
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
656

657
658
659
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
660

661
        for inputs in all_inputs:
662
            output: "GenerateOutput" = self.model.generate(
663
                **self.wrap_device(inputs),
664
665
666
667
668
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
669
                **kwargs,
670
671
            )

672
673
674
            (
                seq_logprobs_lst,
                output_len,
675
            ) = self._hidden_states_to_logprobs(output.hidden_states, num_logprobs)
676
677
678
679
680
681
682

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

684
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
685
686
687
688
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
689

690
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
691
        return self.model.encode(prompts, *args, **kwargs)
692

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

696
697
698
699
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
700
        del self.model
701
        cleanup_dist_env_and_memory()
702

Woosuk Kwon's avatar
Woosuk Kwon committed
703

Cyrus Leung's avatar
Cyrus Leung committed
704
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
705
706
707
708
709
def hf_runner():
    return HfRunner


class VllmRunner:
710
711
    """
    The default value of some arguments have been modified from
712
    {class}`~vllm.LLM` as follows:
713

714
715
716
    - `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.
717
718
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
719
720
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
721
    - `enforce_eager`: Set to `False` to test CUDA graph.
722
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
723
724
725
726

    def __init__(
        self,
        model_name: str,
727
728
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
729
        tokenizer_name: str | None = None,
730
        tokenizer_mode: str = "auto",
731
        trust_remote_code: bool = True,
732
733
        seed: int | None = 0,
        max_model_len: int | None = 1024,
734
        dtype: str = "auto",
735
        disable_log_stats: bool = True,
736
        tensor_parallel_size: int = 1,
737
        block_size: int = 16 if not torch.xpu.is_available() else 64,
738
        enable_chunked_prefill: bool | None = False,
739
        swap_space: int = 4,
740
        enforce_eager: bool | None = False,
741
        # Set this to avoid hanging issue
742
        default_torch_num_threads: int | None = None,
743
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
744
    ) -> None:
745
746
747
748
749
        init_ctx = (
            nullcontext()
            if default_torch_num_threads is None
            else set_default_torch_num_threads(default_torch_num_threads)
        )
750

751
        if not kwargs.get("compilation_config", None):
752
753
754
755
            # 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.
756
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
757

758
759
760
761
762
763
764
765
            # 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
                )

766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
        with init_ctx:
            self.llm = LLM(
                model=model_name,
                runner=runner,
                convert=convert,
                tokenizer=tokenizer_name,
                tokenizer_mode=tokenizer_mode,
                trust_remote_code=trust_remote_code,
                dtype=dtype,
                seed=seed,
                swap_space=swap_space,
                enforce_eager=enforce_eager,
                disable_log_stats=disable_log_stats,
                tensor_parallel_size=tensor_parallel_size,
                max_model_len=max_model_len,
                block_size=block_size,
                enable_chunked_prefill=enable_chunked_prefill,
                **kwargs,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
785

786
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
787
        self,
788
789
790
791
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
792
    ) -> list[dict[str, Any]]:
793
794
795
        if any(
            x is not None and len(x) != len(prompts) for x in [images, videos, audios]
        ):
796
            raise ValueError(
797
798
                "All non-None multimodal inputs must have the same length as prompts"
            )
799

800
        inputs = list[dict[str, Any]]()
801
        for i, prompt in enumerate(prompts):
802
803
804
805
806
807
808
809
810
            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]()
811
812
813
814
815
816
817
            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

818
819
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
820

821
            inputs.append(prompt_dict)
822
823
824
825
826

        return inputs

    def generate(
        self,
827
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
828
        sampling_params: SamplingParams,
829
830
831
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
832
        return_logprobs: bool = False,
833
        **kwargs: Any,
834
    ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
835
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
836

837
838
839
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
840

841
        outputs: list[tuple[list[list[int]], list[str]]] = []
842
        logprobs = []
Woosuk Kwon's avatar
Woosuk Kwon committed
843
844
845
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
846
847
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
848
            req_logprobs = []
849
850
            for sample in req_output.outputs:
                output_str = sample.text
851
                output_ids = list(sample.token_ids)
852
                req_sample_output_ids.append(prompt_ids + output_ids)
853
                req_sample_output_strs.append((prompt_str or "") + output_str)
854
855
                if sample.logprobs:
                    req_logprobs.extend(sample.logprobs)
856
            outputs.append((req_sample_output_ids, req_sample_output_strs))
857
858
            logprobs.append(req_logprobs)
        return outputs if not return_logprobs else (outputs, logprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
859

860
    @staticmethod
861
    def _final_steps_generate_w_logprobs(
862
        req_outputs: list[RequestOutput],
863
        include_prompt_token_ids: bool = False,
864
865
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
866
        for req_output in req_outputs:
867
            assert len(req_output.outputs) > 0
868
869
            for sample in req_output.outputs:
                output_str = sample.text
870
                output_ids = list(sample.token_ids)
871
                output_logprobs = sample.logprobs
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
            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,
                    )
                )

892
893
        return outputs

894
895
    def generate_w_logprobs(
        self,
896
        prompts: list[str],
897
        sampling_params: SamplingParams,
898
899
900
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
901
        include_prompt_token_ids: bool = False,
902
        **kwargs: Any,
903
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
904
905
906
907
908
909
910
        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(
911
            req_outputs, include_prompt_token_ids
912
        )
913
        # Omit prompt logprobs if not required by sampling params
914
915
916
917
918
        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
        )
919

Woosuk Kwon's avatar
Woosuk Kwon committed
920
921
    def generate_greedy(
        self,
922
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
923
        max_tokens: int,
924
925
926
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
927
        **kwargs: Any,
928
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
929
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
930
931
932
933
934
935
936
937
938
        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]
939

940
941
    def generate_greedy_logprobs(
        self,
942
        prompts: list[str],
943
        max_tokens: int,
944
945
946
947
948
949
950
        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,
951
        **kwargs: Any,
952
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
953
954
955
956
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
957
            prompt_logprobs=num_prompt_logprobs,
958
            stop_token_ids=stop_token_ids,
959
960
            stop=stop,
        )
961

962
963
964
965
966
967
968
969
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
970

971
972
973
974
975
976
977
    def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:
        """
        Return the perplexity score associated with generating the prompts

        :param prompts: list of prompts to score
        :return: perplexity score of each prompt
        """
978
979
980
        outputs = self.generate_greedy_logprobs(
            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
        )
981
982
983
984

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
985
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
986
987
988
989
990
991
992
993
994
995
996
997
998
            assert token_datas[0] is None
            token_log_probs = []
            for token_data in token_datas[1:]:
                assert token_data is not None
                assert len(token_data) == 1
                token_log_prob = list(token_data.values())[0].logprob
                token_log_probs.append(token_log_prob)

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

        return perplexities

999
    def generate_beam_search(
1000
        self,
1001
        prompts: list[str],
1002
1003
        beam_width: int,
        max_tokens: int,
1004
1005
1006
1007
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
1008
    ) -> list[tuple[list[list[int]], list[str]]]:
1009
1010
1011
1012
1013
1014
1015
        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,
        )
1016
1017
1018
1019
1020
1021
1022
        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

1023
    def classify(self, prompts: list[str]) -> list[list[float]]:
1024
        req_outputs = self.llm.classify(prompts)
1025
1026
        return [req_output.outputs.probs for req_output in req_outputs]

1027
1028
1029
    def embed(
        self,
        prompts: list[str],
1030
1031
1032
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1033
1034
1035
1036
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1037

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

1041
1042
1043
1044
1045
1046
    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")
1047
1048
        return [req_output.outputs.data for req_output in req_outputs]

1049
1050
1051
1052
    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]

1053
1054
    def score(
        self,
1055
1056
        text_1: list[str] | str,
        text_2: list[str] | str,
1057
1058
        *args,
        **kwargs,
1059
    ) -> list[float]:
1060
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1061
        return [req_output.outputs.score for req_output in req_outputs]
1062

1063
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1064
        return self.llm.apply_model(func)
1065

1066
1067
1068
    def get_llm(self) -> LLM:
        return self.llm

1069
1070
1071
1072
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1073
        del self.llm
1074
        cleanup_dist_env_and_memory()
1075

Woosuk Kwon's avatar
Woosuk Kwon committed
1076

1077
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1078
1079
def vllm_runner():
    return VllmRunner
1080
1081


1082
1083
1084
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1085

1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
    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
1097
1098


1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
@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(),
        }
1181
        config["loggers"]["vllm"]["level"] = level
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194

        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


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

1200
    from vllm.platforms import current_platform
1201

1202
    return current_platform.device_count()
1203
1204
1205


temp_dir = tempfile.gettempdir()
1206
1207
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1208
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1209
1210
1211
1212


@pytest.fixture
def dummy_opt_path():
1213
1214
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1215
1216
1217
1218
1219
        snapshot_download(
            repo_id="facebook/opt-125m",
            local_dir=_dummy_opt_path,
            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
        )
1220
        assert os.path.exists(json_path)
1221
        with open(json_path) as f:
1222
1223
1224
1225
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1226
1227
1228
1229
1230
1231
1232
    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):
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
        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",
            ],
        )
1245
        assert os.path.exists(json_path)
1246
        with open(json_path) as f:
1247
1248
1249
1250
1251
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1252
1253
1254
1255
1256
1257


@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):
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1270
        assert os.path.exists(json_path)
1271
        with open(json_path) as f:
1272
1273
1274
1275
1276
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1277
1278
1279
1280
1281


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1282
1283
1284
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1285
1286
1287
1288
1289
1290
1291
1292
1293


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:
1294
            item.add_marker(skip_optional)
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306


@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")
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353


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
1354
1355
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365

    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(
1366
1367
1368
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
        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:
1396
1397
1398
1399
1400
        image = fetch_image(self.url_for(name))
        # Unwrap MediaWithBytes if present
        if isinstance(image, MediaWithBytes):
            image = image.media
        return image
1401
1402
1403
1404
1405


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
1406
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
    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]
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438


@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()
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467


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

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

    from tests.utils import wait_for_gpu_memory_to_clear

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

    yield

    # Clean up GPU memory after the test
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()