conftest.py 46.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
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.utils import fetch_image
63
from vllm.outputs import RequestOutput
64
from vllm.sampling_params import BeamSearchParams
65
from vllm.transformers_utils.utils import maybe_model_redirect
66
from vllm.utils.collection_utils import is_list_of
67
from vllm.utils.torch_utils import set_default_torch_num_threads
68

69
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
70

71
72
73
_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")]
74
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
75

Cyrus Leung's avatar
Cyrus Leung committed
76
_M = TypeVar("_M")
77

78
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
79
80

PromptImageInput = _PromptMultiModalInput[Image.Image]
81
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
82
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
83

84

85
def _read_prompts(filename: str) -> list[str]:
86
    with open(filename) as f:
87
88
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
89
90


91
class ImageAssetPrompts(TypedDict):
92
93
    stop_sign: str
    cherry_blossom: str
94
95


96
class ImageTestAssets(list[ImageAsset]):
97
    def __init__(self) -> None:
98
99
100
101
102
103
        super().__init__(
            [
                ImageAsset("stop_sign"),
                ImageAsset("cherry_blossom"),
            ]
        )
104

105
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
106
107
108
109
110
111
        """
        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.
        """
112
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
113
114


115
116
class VideoAssetPrompts(TypedDict):
    baby_reading: str
117
118


119
class VideoTestAssets(list[VideoAsset]):
120
    def __init__(self) -> None:
121
122
123
124
125
        super().__init__(
            [
                VideoAsset("baby_reading"),
            ]
        )
126

127
128
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
129
130


131
class AudioAssetPrompts(TypedDict):
132
133
134
135
    mary_had_lamb: str
    winning_call: str


136
class AudioTestAssets(list[AudioAsset]):
137
    def __init__(self) -> None:
138
139
140
141
142
143
        super().__init__(
            [
                AudioAsset("mary_had_lamb"),
                AudioAsset("winning_call"),
            ]
        )
144

145
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
146
147
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

148

149
IMAGE_ASSETS = ImageTestAssets()
150
"""Singleton instance of {class}`ImageTestAssets`."""
151
VIDEO_ASSETS = VideoTestAssets()
152
"""Singleton instance of {class}`VideoTestAssets`."""
153
AUDIO_ASSETS = AudioTestAssets()
154
"""Singleton instance of {class}`AudioTestAssets`."""
155
156


157
158
159
160
161
162
163
@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


164
165
166
167
168
169
170
171
172
173
174
175
@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
176
    cleanup_dist_env_and_memory()
177
178


179
@pytest.fixture()
180
def should_do_global_cleanup_after_test(request) -> bool:
181
182
183
184
    """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.
    """
185

186
    return not request.node.get_closest_marker("skip_global_cleanup")
187
188


189
@pytest.fixture(autouse=True)
190
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
191
    yield
192
    if should_do_global_cleanup_after_test:
193
        cleanup_dist_env_and_memory()
194
195


196
197
198
199
200
201
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
202
@pytest.fixture
203
def example_prompts() -> list[str]:
204
205
    prompts = []
    for filename in _TEST_PROMPTS:
206
        prompts += _read_prompts(filename)
207
208
209
    return prompts


210
211
212
213
214
215
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


216
217
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
218

219
220
221
222
223
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


224
@pytest.fixture
225
def example_long_prompts() -> list[str]:
226
227
    prompts = []
    for filename in _LONG_PROMPTS:
228
        prompts += _read_prompts(filename)
229
    return prompts
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
            model = auto_cls.from_pretrained(
                model_name,
357
                trust_remote_code=trust_remote_code,
358
359
360
                **model_kwargs,
            )

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

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

            self.model = model
374

375
376
377
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
378
                dtype=dtype,
379
                trust_remote_code=trust_remote_code,
380
            )
381

382
383
384
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
385

386
387
        self.processor = AutoProcessor.from_pretrained(
            model_name,
388
            dtype=dtype,
389
            trust_remote_code=trust_remote_code,
390
        )
391
392
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
393

394
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
395
        self,
396
397
398
399
400
        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]]:
401
        if images is not None:
402
            assert len(prompts) == len(images)
403

404
405
406
407
408
409
        if videos is not None:
            assert len(prompts) == len(videos)

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

410
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
411
        for i, prompt in enumerate(prompts):
412
413
414
415
416
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
            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)
450
451
452

        return all_inputs

453
454
455
456
457
458
459
460
461
    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

462
    def classify(self, prompts: list[str]) -> list[str]:
463
464
465
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
466
467
        problem_type = getattr(self.config, "problem_type", "")

468
469
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
470
471
472
473
474
475
            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()
476
477
478
479
            outputs.append(logits)

        return outputs

480
481
    def generate(
        self,
482
483
484
485
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
486
        **kwargs: Any,
487
    ) -> list[tuple[list[list[int]], list[str]]]:
488
489
490
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
491

492
        outputs: list[tuple[list[list[int]], list[str]]] = []
493
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
494
            output_ids = self.model.generate(
495
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
496
497
498
                use_cache=True,
                **kwargs,
            )
499
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
500
501
502
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
503
504
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
505
506
507
508
509
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
510
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
511
        max_tokens: int,
512
513
514
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
515
        **kwargs: Any,
516
    ) -> list[tuple[list[int], str]]:
517
518
519
520
521
522
523
524
525
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
526

527
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
528
529
530

    def generate_beam_search(
        self,
531
        prompts: list[str],
532
533
        beam_width: int,
        max_tokens: int,
534
535
536
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
537
    ) -> list[tuple[list[list[int]], list[str]]]:
538
539
540
541
542
543
544
545
546
547
        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,
        )
548

549
550
551
552
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
553
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
554
555
556
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
557

558
559
    def generate_greedy_logprobs(
        self,
560
        prompts: list[str],
561
        max_tokens: int,
562
563
564
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
565
        **kwargs: Any,
566
    ) -> list[list[torch.Tensor]]:
567
568
569
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
570

571
        all_logprobs: list[list[torch.Tensor]] = []
572
        for inputs in all_inputs:
573
            output = self.model.generate(
574
                **self.wrap_device(inputs),
575
576
577
578
579
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
580
                **kwargs,
581
            )
582
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
583
584
585
            all_logprobs.append(seq_logprobs)
        return all_logprobs

586
    def _hidden_states_to_seq_logprobs(
587
        self,
588
589
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
590
591
        output_embeddings = self.model.get_output_embeddings()

592
        seq_logprobs: list[torch.Tensor] = []
593
594
595
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
596
597
598
599
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
600
                output_embeddings.weight.t(),
601
            )
602
603
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
604
605
606
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

607
608
609
610
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
611
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
612
        num_logprobs: int | None,
613
    ) -> tuple[list[dict[int, float]], int]:
614
615
616
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

617
        # convert to dict
618
        seq_logprobs_lst: list[dict[int, float]] = []
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
        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,
        )

636
637
    def generate_greedy_logprobs_limit(
        self,
638
        prompts: list[str],
639
        max_tokens: int,
640
641
642
643
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
644
        **kwargs: Any,
645
    ) -> list[TokensTextLogprobs]:
646
647
648
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
649

650
651
652
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
653

654
        for inputs in all_inputs:
655
            output = self.model.generate(
656
                **self.wrap_device(inputs),
657
658
659
660
661
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
662
                **kwargs,
663
664
            )

665
666
667
            (
                seq_logprobs_lst,
                output_len,
668
            ) = self._hidden_states_to_logprobs(output.hidden_states, num_logprobs)
669
670
671
672
673
674
675

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

677
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
678
679
680
681
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
682

683
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
684
        return self.model.encode(prompts, *args, **kwargs)
685

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

689
690
691
692
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
693
        del self.model
694
        cleanup_dist_env_and_memory()
695

Woosuk Kwon's avatar
Woosuk Kwon committed
696

Cyrus Leung's avatar
Cyrus Leung committed
697
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
698
699
700
701
702
def hf_runner():
    return HfRunner


class VllmRunner:
703
704
    """
    The default value of some arguments have been modified from
705
    {class}`~vllm.LLM` as follows:
706

707
708
709
    - `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.
710
711
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
712
713
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
714
    - `enforce_eager`: Set to `False` to test CUDA graph.
715
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
716
717
718
719

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

744
        if not kwargs.get("compilation_config", None):
745
746
747
748
            # 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.
749
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
750

751
752
753
754
755
756
757
758
            # 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
                )

759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
        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
778

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

793
        inputs = list[dict[str, Any]]()
794
        for i, prompt in enumerate(prompts):
795
796
797
798
799
800
801
802
803
            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]()
804
805
806
807
808
809
810
            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

811
812
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
813

814
            inputs.append(prompt_dict)
815
816
817
818
819

        return inputs

    def generate(
        self,
820
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
821
        sampling_params: SamplingParams,
822
823
824
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
825
        return_logprobs: bool = False,
826
        **kwargs: Any,
827
    ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
828
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
829

830
831
832
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
833

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

853
    @staticmethod
854
    def _final_steps_generate_w_logprobs(
855
856
857
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
858
        for req_output in req_outputs:
859
            assert len(req_output.outputs) > 0
860
861
            for sample in req_output.outputs:
                output_str = sample.text
862
                output_ids = list(sample.token_ids)
863
                output_logprobs = sample.logprobs
864
865
866
            outputs.append(
                (output_ids, output_str, output_logprobs, req_output.prompt_logprobs)
            )
867
868
        return outputs

869
870
    def generate_w_logprobs(
        self,
871
        prompts: list[str],
872
        sampling_params: SamplingParams,
873
874
875
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
876
        **kwargs: Any,
877
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
878
879
880
881
882
883
884
885
886
        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(
            req_outputs
        )
887
        # Omit prompt logprobs if not required by sampling params
888
889
890
891
892
        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
        )
893

Woosuk Kwon's avatar
Woosuk Kwon committed
894
895
    def generate_greedy(
        self,
896
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
897
        max_tokens: int,
898
899
900
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
901
        **kwargs: Any,
902
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
903
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
904
905
906
907
908
909
910
911
912
        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]
913

914
915
    def generate_greedy_logprobs(
        self,
916
        prompts: list[str],
917
        max_tokens: int,
918
919
920
921
922
923
924
        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,
925
        **kwargs: Any,
926
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
927
928
929
930
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
931
            prompt_logprobs=num_prompt_logprobs,
932
            stop_token_ids=stop_token_ids,
933
934
            stop=stop,
        )
935

936
937
938
939
940
941
942
943
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
944

945
946
947
948
949
950
951
    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
        """
952
953
954
        outputs = self.generate_greedy_logprobs(
            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
        )
955
956
957
958

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
959
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
960
961
962
963
964
965
966
967
968
969
970
971
972
            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

973
    def generate_beam_search(
974
        self,
975
        prompts: list[str],
976
977
        beam_width: int,
        max_tokens: int,
978
979
980
981
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
982
    ) -> list[tuple[list[list[int]], list[str]]]:
983
984
985
986
987
988
989
        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,
        )
990
991
992
993
994
995
996
        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

997
    def classify(self, prompts: list[str]) -> list[list[float]]:
998
        req_outputs = self.llm.classify(prompts)
999
1000
        return [req_output.outputs.probs for req_output in req_outputs]

1001
1002
1003
    def embed(
        self,
        prompts: list[str],
1004
1005
1006
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1007
1008
1009
1010
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1011

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

1015
1016
1017
1018
1019
1020
    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")
1021
1022
        return [req_output.outputs.data for req_output in req_outputs]

1023
1024
1025
1026
    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]

1027
1028
    def score(
        self,
1029
1030
        text_1: list[str] | str,
        text_2: list[str] | str,
1031
1032
        *args,
        **kwargs,
1033
    ) -> list[float]:
1034
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1035
        return [req_output.outputs.score for req_output in req_outputs]
1036

1037
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1038
        return self.llm.apply_model(func)
1039

1040
1041
1042
    def get_llm(self) -> LLM:
        return self.llm

1043
1044
1045
1046
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1047
        del self.llm
1048
        cleanup_dist_env_and_memory()
1049

Woosuk Kwon's avatar
Woosuk Kwon committed
1050

1051
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1052
1053
def vllm_runner():
    return VllmRunner
1054
1055


1056
1057
1058
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1059

1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
    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
1071
1072


1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
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
@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(),
        }

        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


1168
1169
1170
1171
1172
@pytest.fixture(scope="session")
def num_gpus_available():
    """Get number of GPUs without initializing the CUDA context
    in current process."""

1173
    from vllm.platforms import current_platform
1174

1175
    return current_platform.device_count()
1176
1177
1178


temp_dir = tempfile.gettempdir()
1179
1180
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1181
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1182
1183
1184
1185


@pytest.fixture
def dummy_opt_path():
1186
1187
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1188
1189
1190
1191
1192
        snapshot_download(
            repo_id="facebook/opt-125m",
            local_dir=_dummy_opt_path,
            ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"],
        )
1193
        assert os.path.exists(json_path)
1194
        with open(json_path) as f:
1195
1196
1197
1198
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1199
1200
1201
1202
1203
1204
1205
    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):
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
        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",
            ],
        )
1218
        assert os.path.exists(json_path)
1219
        with open(json_path) as f:
1220
1221
1222
1223
1224
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1225
1226
1227
1228
1229
1230


@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):
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1243
        assert os.path.exists(json_path)
1244
        with open(json_path) as f:
1245
1246
1247
1248
1249
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1250
1251
1252
1253
1254


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1255
1256
1257
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1258
1259
1260
1261
1262
1263
1264
1265
1266


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:
1267
            item.add_marker(skip_optional)
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279


@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")
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326


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
1327
1328
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338

    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(
1339
1340
1341
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
        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:
        return fetch_image(self.url_for(name))


@pytest.fixture(scope="session")
def local_asset_server() -> Generator[LocalAssetServer, None, None]:
    """
1375
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
    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]
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407


@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()