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

from tblib import pickling_support

9
10
# ruff: noqa

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

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

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

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

66
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
67

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

Cyrus Leung's avatar
Cyrus Leung committed
73
_M = TypeVar("_M")
74

75
_PromptMultiModalInput = list[_M] | list[list[_M]]
Cyrus Leung's avatar
Cyrus Leung committed
76
77

PromptImageInput = _PromptMultiModalInput[Image.Image]
78
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
79
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
80

81

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


88
class ImageAssetPrompts(TypedDict):
89
90
    stop_sign: str
    cherry_blossom: str
91
92


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

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


112
113
class VideoAssetPrompts(TypedDict):
    baby_reading: str
114
115


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

124
125
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
126
127


128
class AudioAssetPrompts(TypedDict):
129
130
131
132
    mary_had_lamb: str
    winning_call: str


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

142
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
143
144
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

145

146
IMAGE_ASSETS = ImageTestAssets()
147
"""Singleton instance of {class}`ImageTestAssets`."""
148
VIDEO_ASSETS = VideoTestAssets()
149
"""Singleton instance of {class}`VideoTestAssets`."""
150
AUDIO_ASSETS = AudioTestAssets()
151
"""Singleton instance of {class}`AudioTestAssets`."""
152
153


154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
@pytest.fixture(scope="function", autouse=True)
def cleanup_VLLM_USE_V1(monkeypatch):
    """
    The V1 oracle sets "VLLM_USE_V1" during loading. This means
    that each invocation of a test change the env variable.

    If we touch "VLLM_USE_V1" with monkeypatch, then any changes
    made during the test run by vLLM will be cleaned up.

    This fixture is used by every test.
    """

    # If VLLM_USE_V1 is not set, set then delete. This will
    # cause monkeypatch to clean up VLLM_USE_V1 upon exit
    # if VLLM modifies the value of envs.VLLM_USE_V1.
    if "VLLM_USE_V1" not in os.environ:
        monkeypatch.setenv("VLLM_USE_V1", "")
        monkeypatch.delenv("VLLM_USE_V1")


174
175
176
177
178
179
180
@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


181
182
183
184
185
186
187
188
189
190
191
192
@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
193
    cleanup_dist_env_and_memory()
194
195


196
@pytest.fixture()
197
def should_do_global_cleanup_after_test(request) -> bool:
198
199
200
201
    """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.
    """
202

203
    return not request.node.get_closest_marker("skip_global_cleanup")
204
205


206
@pytest.fixture(autouse=True)
207
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
208
    yield
209
    if should_do_global_cleanup_after_test:
210
        cleanup_dist_env_and_memory()
211
212


213
214
215
216
217
218
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
219
@pytest.fixture
220
def example_prompts() -> list[str]:
221
222
    prompts = []
    for filename in _TEST_PROMPTS:
223
        prompts += _read_prompts(filename)
224
225
226
    return prompts


227
228
229
230
231
232
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


233
234
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
235

236
237
238
239
240
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


241
@pytest.fixture
242
def example_long_prompts() -> list[str]:
243
244
    prompts = []
    for filename in _LONG_PROMPTS:
245
        prompts += _read_prompts(filename)
246
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
247
248


249
@pytest.fixture(scope="session")
250
def image_assets() -> ImageTestAssets:
251
252
253
    return IMAGE_ASSETS


254
@pytest.fixture(scope="session")
255
def video_assets() -> VideoTestAssets:
256
257
258
    return VIDEO_ASSETS


259
@pytest.fixture(scope="session")
260
def audio_assets() -> AudioTestAssets:
261
262
263
    return AUDIO_ASSETS


264
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
265
_R = TypeVar("_R")
266

Woosuk Kwon's avatar
Woosuk Kwon committed
267
268

class HfRunner:
269
    def get_default_device(self):
270
        from vllm.platforms import current_platform
271

272
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
273

274
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
275
        if x is None or isinstance(x, (bool,)):
276
277
            return x

278
        if device is None:
279
            device = self.device
280

281
282
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
283

284
285
286
287
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
288

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

        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",
        *,
326
        model_kwargs: dict[str, Any] | None = None,
327
328
329
330
331
        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
332
    ) -> None:
333
        model_name = maybe_model_redirect(model_name)
334
        self.model_name = model_name
335

336
337
        self.config = AutoConfig.from_pretrained(
            model_name,
338
            trust_remote_code=trust_remote_code,
339
340
        )
        self.device = self.get_default_device()
341
        self.dtype = dtype = _get_and_verify_dtype(
342
343
344
345
346
            self.model_name,
            self.config,
            dtype=dtype,
            is_pooling_model=is_sentence_transformer or is_cross_encoder,
        )
347
348

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

351
        if is_sentence_transformer:
352
353
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
354
355
356
357
358

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
359
                trust_remote_code=trust_remote_code,
360
            )
361
362
363
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
364
365
366
367
368

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
369
                trust_remote_code=trust_remote_code,
370
            )
371
        else:
372
373
            model = auto_cls.from_pretrained(
                model_name,
374
                trust_remote_code=trust_remote_code,
375
376
377
                **model_kwargs,
            )

378
            # in case some unquantized custom models are not in same dtype
379
380
381
            if getattr(model, "quantization_method", None) is None and any(
                p.dtype != self.dtype for p in model.parameters()
            ):
382
383
                model = model.to(dtype=self.dtype)

384
385
386
387
            if (
                getattr(model, "quantization_method", None) != "bitsandbytes"
                and len({p.device for p in model.parameters()}) < 2
            ):
388
                model = model.to(device=self.device)
389
390

            self.model = model
391

392
393
394
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
395
                dtype=dtype,
396
                trust_remote_code=trust_remote_code,
397
            )
398

399
400
401
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
402

403
404
        self.processor = AutoProcessor.from_pretrained(
            model_name,
405
            dtype=dtype,
406
            trust_remote_code=trust_remote_code,
407
        )
408
409
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
410

411
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
412
        self,
413
414
415
416
417
        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]]:
418
        if images is not None:
419
            assert len(prompts) == len(images)
420

421
422
423
424
425
426
        if videos is not None:
            assert len(prompts) == len(videos)

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

427
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
428
        for i, prompt in enumerate(prompts):
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
455
456
457
458
459
460
461
462
463
464
465
466
            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)
467
468
469

        return all_inputs

470
471
472
473
474
475
476
477
478
    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

479
    def classify(self, prompts: list[str]) -> list[str]:
480
481
482
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
483
484
        problem_type = getattr(self.config, "problem_type", "")

485
486
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
487
488
489
490
491
492
            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()
493
494
495
496
            outputs.append(logits)

        return outputs

497
498
    def generate(
        self,
499
500
501
502
        prompts: list[str] | list[list[int]],
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
503
        **kwargs: Any,
504
    ) -> list[tuple[list[list[int]], list[str]]]:
505
506
507
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
508

509
        outputs: list[tuple[list[list[int]], list[str]]] = []
510
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
511
            output_ids = self.model.generate(
512
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
513
514
515
                use_cache=True,
                **kwargs,
            )
516
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
517
518
519
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
520
521
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
522
523
524
525
526
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
527
        prompts: list[str] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
528
        max_tokens: int,
529
530
531
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
532
        **kwargs: Any,
533
    ) -> list[tuple[list[int], str]]:
534
535
536
537
538
539
540
541
542
        outputs = self.generate(
            prompts,
            do_sample=False,
            max_new_tokens=max_tokens,
            images=images,
            videos=videos,
            audios=audios,
            **kwargs,
        )
543

544
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
545
546
547

    def generate_beam_search(
        self,
548
        prompts: list[str],
549
550
        beam_width: int,
        max_tokens: int,
551
552
553
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
554
    ) -> list[tuple[list[list[int]], list[str]]]:
555
556
557
558
559
560
561
562
563
564
        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,
        )
565

566
567
568
569
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
570
                    x for x in output_ids[j] if x != self.tokenizer.pad_token_id
571
572
573
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
574

575
576
    def generate_greedy_logprobs(
        self,
577
        prompts: list[str],
578
        max_tokens: int,
579
580
581
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
582
        **kwargs: Any,
583
    ) -> list[list[torch.Tensor]]:
584
585
586
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
587

588
        all_logprobs: list[list[torch.Tensor]] = []
589
        for inputs in all_inputs:
590
            output = self.model.generate(
591
                **self.wrap_device(inputs),
592
593
594
595
596
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
597
                **kwargs,
598
            )
599
            seq_logprobs = self._hidden_states_to_seq_logprobs(output.hidden_states)
600
601
602
            all_logprobs.append(seq_logprobs)
        return all_logprobs

603
    def _hidden_states_to_seq_logprobs(
604
        self,
605
606
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
    ) -> list[torch.Tensor]:
607
608
        output_embeddings = self.model.get_output_embeddings()

609
        seq_logprobs: list[torch.Tensor] = []
610
611
612
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
613
614
615
616
                last_hidden_states.to(
                    device=output_embeddings.weight.device,
                    dtype=output_embeddings.weight.dtype,
                ),
617
                output_embeddings.weight.t(),
618
            )
619
620
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
621
622
623
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

624
625
626
627
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
628
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
629
        num_logprobs: int | None,
630
    ) -> tuple[list[dict[int, float]], int]:
631
632
633
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

634
        # convert to dict
635
        seq_logprobs_lst: list[dict[int, float]] = []
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
        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,
        )

653
654
    def generate_greedy_logprobs_limit(
        self,
655
        prompts: list[str],
656
        max_tokens: int,
657
658
659
660
        num_logprobs: int | None,
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
661
        **kwargs: Any,
662
    ) -> list[TokensTextLogprobs]:
663
664
665
        all_inputs = self.get_inputs(
            prompts, images=images, videos=videos, audios=audios
        )
666

667
668
669
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
670

671
        for inputs in all_inputs:
672
            output = self.model.generate(
673
                **self.wrap_device(inputs),
674
675
676
677
678
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
679
                **kwargs,
680
681
            )

682
683
684
            (
                seq_logprobs_lst,
                output_len,
685
            ) = self._hidden_states_to_logprobs(output.hidden_states, num_logprobs)
686
687
688
689
690
691
692

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

694
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
695
696
697
698
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
699

700
    def encode(self, prompts: list[str], *args, **kwargs) -> list[list[torch.Tensor]]:
701
        return self.model.encode(prompts, *args, **kwargs)
702

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

706
707
708
709
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
710
        del self.model
711
        cleanup_dist_env_and_memory()
712

Woosuk Kwon's avatar
Woosuk Kwon committed
713

Cyrus Leung's avatar
Cyrus Leung committed
714
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
715
716
717
718
719
def hf_runner():
    return HfRunner


class VllmRunner:
720
721
    """
    The default value of some arguments have been modified from
722
    {class}`~vllm.LLM` as follows:
723

724
725
726
    - `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.
727
728
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
729
730
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
731
    - `enforce_eager`: Set to `False` to test CUDA graph.
732
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
733
734
735
736

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

761
        if not kwargs.get("compilation_config", None):
762
763
764
765
            # 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.
766
            kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
767

768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
        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
787

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

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

820
821
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
822

823
            inputs.append(prompt_dict)
824
825
826
827
828

        return inputs

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

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

842
        outputs: list[tuple[list[list[int]], list[str]]] = []
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
849
            for sample in req_output.outputs:
                output_str = sample.text
850
                output_ids = list(sample.token_ids)
851
                req_sample_output_ids.append(prompt_ids + output_ids)
852
                req_sample_output_strs.append((prompt_str or "") + output_str)
853
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
854
855
        return outputs

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

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

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

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

939
940
941
942
943
944
945
946
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
947

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

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

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

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

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

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

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

1026
1027
1028
1029
    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]

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

1040
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1041
        return self.llm.apply_model(func)
1042

1043
1044
1045
    def get_llm(self) -> LLM:
        return self.llm

1046
1047
1048
1049
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1050
        del self.llm
1051
        cleanup_dist_env_and_memory()
1052

Woosuk Kwon's avatar
Woosuk Kwon committed
1053

1054
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1055
1056
def vllm_runner():
    return VllmRunner
1057
1058


1059
1060
1061
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1062

1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
    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
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
1168
1169
1170
@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


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

1176
    from vllm.platforms import current_platform
1177

1178
    return current_platform.device_count()
1179
1180
1181


temp_dir = tempfile.gettempdir()
1182
1183
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1184
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1185
1186
1187
1188


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


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


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


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


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


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

    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(
1342
1343
1344
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
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
1375
1376
1377
        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]:
    """
1378
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
    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]