conftest.py 46.3 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
164
165
166
167
168
169
170
171
172
173
174
175
176
@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")


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


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


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

206
    return not request.node.get_closest_marker("skip_global_cleanup")
207
208


209
@pytest.fixture(autouse=True)
210
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
211
    yield
212
    if should_do_global_cleanup_after_test:
213
        cleanup_dist_env_and_memory()
214
215


216
217
218
219
220
221
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


230
231
232
233
234
235
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


236
237
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
238

239
240
241
242
243
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


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


252
@pytest.fixture(scope="session")
253
def image_assets() -> ImageTestAssets:
254
255
256
    return IMAGE_ASSETS


257
@pytest.fixture(scope="session")
258
def video_assets() -> VideoTestAssets:
259
260
261
    return VIDEO_ASSETS


262
@pytest.fixture(scope="session")
263
def audio_assets() -> AudioTestAssets:
264
265
266
    return AUDIO_ASSETS


267
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
268
_R = TypeVar("_R")
269

Woosuk Kwon's avatar
Woosuk Kwon committed
270
271

class HfRunner:
272
    def get_default_device(self):
273
        from vllm.platforms import current_platform
274

275
        return "cpu" if current_platform.is_cpu() else current_platform.device_type
276

277
    def wrap_device(self, x: _T, device: str | None = None) -> _T:
278
        if x is None or isinstance(x, (bool,)):
279
280
            return x

281
        if device is None:
282
            device = self.device
283

284
285
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
286

287
288
289
290
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
291

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

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

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

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

354
        if is_sentence_transformer:
355
356
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
357
358
359
360
361

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

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

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

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

            self.model = model
394

395
396
397
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
398
                dtype=dtype,
399
                trust_remote_code=trust_remote_code,
400
            )
401

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

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

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

424
425
426
427
428
429
        if videos is not None:
            assert len(prompts) == len(videos)

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

430
        all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []
431
        for i, prompt in enumerate(prompts):
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
467
468
469
            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)
470
471
472

        return all_inputs

473
474
475
476
477
478
479
480
481
    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

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

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

        return outputs

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

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

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

547
        return [(output_ids[0], output_str[0]) for output_ids, output_str in outputs]
548
549
550

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

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

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

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

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

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

627
628
629
630
        return seq_logprobs

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

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

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

670
671
672
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
673

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

685
686
687
            (
                seq_logprobs_lst,
                output_len,
688
            ) = self._hidden_states_to_logprobs(output.hidden_states, num_logprobs)
689
690
691
692
693
694
695

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

697
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
698
699
700
701
        return [
            (output_ids, output_str, output_logprobs)
            for output_ids, output_str, output_logprobs in outputs
        ]
702

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

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

709
710
711
712
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
713
        del self.model
714
        cleanup_dist_env_and_memory()
715

Woosuk Kwon's avatar
Woosuk Kwon committed
716

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


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

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

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

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

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

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

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

823
824
            if multi_modal_data:
                prompt_dict["multi_modal_data"] = multi_modal_data
825

826
            inputs.append(prompt_dict)
827
828
829
830
831

        return inputs

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

842
843
844
        req_outputs = self.llm.generate(
            inputs, sampling_params=sampling_params, **kwargs
        )
845

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

865
    @staticmethod
866
    def _final_steps_generate_w_logprobs(
867
868
869
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
870
        for req_output in req_outputs:
871
            assert len(req_output.outputs) > 0
872
873
            for sample in req_output.outputs:
                output_str = sample.text
874
                output_ids = list(sample.token_ids)
875
                output_logprobs = sample.logprobs
876
877
878
            outputs.append(
                (output_ids, output_str, output_logprobs, req_output.prompt_logprobs)
            )
879
880
        return outputs

881
882
    def generate_w_logprobs(
        self,
883
        prompts: list[str],
884
        sampling_params: SamplingParams,
885
886
887
        images: PromptImageInput | None = None,
        audios: PromptAudioInput | None = None,
        videos: PromptVideoInput | None = None,
888
        **kwargs: Any,
889
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
890
891
892
893
894
895
896
897
898
        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
        )
899
        # Omit prompt logprobs if not required by sampling params
900
901
902
903
904
        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
        )
905

Woosuk Kwon's avatar
Woosuk Kwon committed
906
907
    def generate_greedy(
        self,
908
        prompts: list[str] | list[torch.Tensor] | list[list[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
909
        max_tokens: int,
910
911
912
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
913
        **kwargs: Any,
914
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
915
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
916
917
918
919
920
921
922
923
924
        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]
925

926
927
    def generate_greedy_logprobs(
        self,
928
        prompts: list[str],
929
        max_tokens: int,
930
931
932
933
934
935
936
        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,
937
        **kwargs: Any,
938
    ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:
939
940
941
942
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
943
            prompt_logprobs=num_prompt_logprobs,
944
            stop_token_ids=stop_token_ids,
945
946
            stop=stop,
        )
947

948
949
950
951
952
953
954
955
        return self.generate_w_logprobs(
            prompts,
            greedy_logprobs_params,
            images=images,
            audios=audios,
            videos=videos,
            **kwargs,
        )
956

957
958
959
960
961
962
963
    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
        """
964
965
966
        outputs = self.generate_greedy_logprobs(
            prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
        )
967
968
969
970

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
971
            token_datas = cast(list[dict[int, Logprob] | None], output[3])
972
973
974
975
976
977
978
979
980
981
982
983
984
            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

985
    def generate_beam_search(
986
        self,
987
        prompts: list[str],
988
989
        beam_width: int,
        max_tokens: int,
990
991
992
993
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
        concurrency_limit: int | None = None,
994
    ) -> list[tuple[list[list[int]], list[str]]]:
995
996
997
998
999
1000
1001
        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,
        )
1002
1003
1004
1005
1006
1007
1008
        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

1009
    def classify(self, prompts: list[str]) -> list[list[float]]:
1010
        req_outputs = self.llm.classify(prompts)
1011
1012
        return [req_output.outputs.probs for req_output in req_outputs]

1013
1014
1015
    def embed(
        self,
        prompts: list[str],
1016
1017
1018
        images: PromptImageInput | None = None,
        videos: PromptVideoInput | None = None,
        audios: PromptAudioInput | None = None,
1019
1020
1021
1022
        *args,
        **kwargs,
    ) -> list[list[float]]:
        inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
Cyrus Leung's avatar
Cyrus Leung committed
1023

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

1027
1028
1029
1030
1031
1032
    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")
1033
1034
        return [req_output.outputs.data for req_output in req_outputs]

1035
1036
1037
1038
    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]

1039
1040
    def score(
        self,
1041
1042
        text_1: list[str] | str,
        text_2: list[str] | str,
1043
1044
        *args,
        **kwargs,
1045
    ) -> list[float]:
1046
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1047
        return [req_output.outputs.score for req_output in req_outputs]
1048

1049
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1050
        return self.llm.apply_model(func)
1051

1052
1053
1054
    def get_llm(self) -> LLM:
        return self.llm

1055
1056
1057
1058
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1059
        del self.llm
1060
        cleanup_dist_env_and_memory()
1061

Woosuk Kwon's avatar
Woosuk Kwon committed
1062

1063
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1064
1065
def vllm_runner():
    return VllmRunner
1066
1067


1068
1069
1070
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
1071

1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
    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
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
1171
1172
1173
1174
1175
1176
1177
1178
1179
@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


1180
1181
1182
1183
1184
@pytest.fixture(scope="session")
def num_gpus_available():
    """Get number of GPUs without initializing the CUDA context
    in current process."""

1185
    from vllm.platforms import current_platform
1186

1187
    return current_platform.device_count()
1188
1189
1190


temp_dir = tempfile.gettempdir()
1191
1192
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1193
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1194
1195
1196
1197


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


@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):
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
        snapshot_download(
            repo_id="BAAI/bge-multilingual-gemma2",
            local_dir=_dummy_gemma2_embedding_path,
            ignore_patterns=[
                "*.bin",
                "*.bin.index.json",
                "*.pt",
                "*.h5",
                "*.msgpack",
                "*.safetensors",
            ],
        )
1255
        assert os.path.exists(json_path)
1256
        with open(json_path) as f:
1257
1258
1259
1260
1261
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1262
1263
1264
1265
1266


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
1267
1268
1269
    parser.addoption(
        "--optional", action="store_true", default=False, help="run optional test"
    )
1270
1271
1272
1273
1274
1275
1276
1277
1278


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:
1279
            item.add_marker(skip_optional)
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291


@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")
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
1330
1331
1332
1333
1334
1335
1336
1337
1338


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
1339
1340
    server: http.server.ThreadingHTTPServer | None
    thread: threading.Thread | None
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350

    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(
1351
1352
1353
            (self.address, self.port), AssetHandler
        )
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
        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]:
    """
1387
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port.
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
    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]