conftest.py 47.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import http.server
4
import json
5
import math
6
import mimetypes
7
import os
8
import socket
9
import tempfile
10
11
import threading
from collections.abc import Generator
12
from enum import Enum
13
from typing import Any, Callable, Optional, TypedDict, TypeVar, Union, cast
Woosuk Kwon's avatar
Woosuk Kwon committed
14

15
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
16
17
import pytest
import torch
18
import torch.nn as nn
19
import torch.nn.functional as F
20
from huggingface_hub import snapshot_download
21
from PIL import Image
22
23
from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
                          BatchEncoding, BatchFeature)
24
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
25

26
27
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
28
from vllm import LLM, SamplingParams
29
from vllm.assets.audio import AudioAsset
30
from vllm.assets.image import ImageAsset
31
from vllm.assets.video import VideoAsset
32
from vllm.config import ConvertOption, RunnerOption, _get_and_verify_dtype
33
from vllm.connections import global_http_connection
34
from vllm.distributed import (cleanup_dist_env_and_memory,
35
36
                              init_distributed_environment,
                              initialize_model_parallel)
37
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
38
                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
39
from vllm.logger import init_logger
40
from vllm.multimodal.utils import fetch_image
41
from vllm.outputs import RequestOutput
42
from vllm.sampling_params import BeamSearchParams
43
from vllm.sequence import Logprob
44
from vllm.transformers_utils.utils import maybe_model_redirect
45

46
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
47

48
49
50
_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")]
51
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
52

Cyrus Leung's avatar
Cyrus Leung committed
53
_M = TypeVar("_M")
54

55
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
56
57

PromptImageInput = _PromptMultiModalInput[Image.Image]
58
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
59
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
60

61

62
def _read_prompts(filename: str) -> list[str]:
63
    with open(filename) as f:
64
65
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
66
67


68
class ImageAssetPrompts(TypedDict):
69
70
    stop_sign: str
    cherry_blossom: str
71
72


73
class ImageTestAssets(list[ImageAsset]):
74
75

    def __init__(self) -> None:
76
77
78
79
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
80

81
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
82
83
84
85
86
87
        """
        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.
        """
88
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
89
90


91
92
class VideoAssetPrompts(TypedDict):
    baby_reading: str
93
94


95
class VideoTestAssets(list[VideoAsset]):
96
97
98

    def __init__(self) -> None:
        super().__init__([
99
            VideoAsset("baby_reading"),
100
101
        ])

102
103
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
104
105


106
class AudioAssetPrompts(TypedDict):
107
108
109
110
    mary_had_lamb: str
    winning_call: str


111
class AudioTestAssets(list[AudioAsset]):
112
113
114
115
116
117
118

    def __init__(self) -> None:
        super().__init__([
            AudioAsset("mary_had_lamb"),
            AudioAsset("winning_call"),
        ])

119
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
120
121
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

122

123
IMAGE_ASSETS = ImageTestAssets()
124
"""Singleton instance of {class}`ImageTestAssets`."""
125
VIDEO_ASSETS = VideoTestAssets()
126
"""Singleton instance of {class}`VideoTestAssets`."""
127
AUDIO_ASSETS = AudioTestAssets()
128
"""Singleton instance of {class}`AudioTestAssets`."""
129
130


131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@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")


Joe Runde's avatar
Joe Runde committed
151
@pytest.fixture(params=[True, False])
152
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
153
154
155
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
156
    skip_v0 = request.node.get_closest_marker("skip_v0")
Joe Runde's avatar
Joe Runde committed
157
158
159
160
161
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
162
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
163
    else:
164
165
        if skip_v0:
            pytest.skip("Skipping test on vllm V0")
166
167
168
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
169
170


171
172
173
174
175
176
177
@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


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


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

200
    return not request.node.get_closest_marker("skip_global_cleanup")
201
202


203
@pytest.fixture(autouse=True)
204
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
205
    yield
206
    if should_do_global_cleanup_after_test:
207
        cleanup_dist_env_and_memory()
208
209


210
211
212
213
214
215
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


224
225
226
227
228
229
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


230
231
232
233
234
235
236
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


237
@pytest.fixture
238
def example_encoder_decoder_prompts(
239
) -> dict[DecoderPromptType, list[ExplicitEncoderDecoderPrompt]]:
240
241
242
243
244
245
    '''
    Returns an encoder prompt list and a decoder prompt list, wherein each pair
    of same-index entries in both lists corresponds to an (encoder prompt,
    decoder prompt) tuple.

    Returns:
246

247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
    * Encoder prompt list
    * Decoder prompt list (reverse of encoder prompt list)
    '''

    encoder_prompts = []
    for filename in _TEST_PROMPTS:
        encoder_prompts += _read_prompts(filename)

    custom_decoder_prompts = encoder_prompts[::-1]
    empty_str_decoder_prompts = [""] * len(encoder_prompts)
    none_decoder_prompts = [None] * len(encoder_prompts)

    # NONE decoder prompt type
    return {
        DecoderPromptType.NONE:
262
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
263
        DecoderPromptType.EMPTY_STR:
264
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
265
        DecoderPromptType.CUSTOM:
266
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
267
268
269
    }


270
@pytest.fixture
271
def example_long_prompts() -> list[str]:
272
273
    prompts = []
    for filename in _LONG_PROMPTS:
274
        prompts += _read_prompts(filename)
275
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
276
277


278
@pytest.fixture(scope="session")
279
def image_assets() -> ImageTestAssets:
280
281
282
    return IMAGE_ASSETS


283
@pytest.fixture(scope="session")
284
def video_assets() -> VideoTestAssets:
285
286
287
    return VIDEO_ASSETS


288
@pytest.fixture(scope="session")
289
def audio_assets() -> AudioTestAssets:
290
291
292
    return AUDIO_ASSETS


293
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
294
_R = TypeVar("_R")
295

Woosuk Kwon's avatar
Woosuk Kwon committed
296
297
298

class HfRunner:

299
    def get_default_device(self):
300
        from vllm.platforms import current_platform
301

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

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

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

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

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

        return x.to(device)
319

Woosuk Kwon's avatar
Woosuk Kwon committed
320
321
322
    def __init__(
        self,
        model_name: str,
323
        dtype: str = "auto",
324
        *,
325
        model_kwargs: Optional[dict[str, Any]] = None,
326
        trust_remote_code: bool = True,
327
        is_sentence_transformer: bool = False,
328
        is_cross_encoder: bool = False,
329
        skip_tokenizer_init: bool = False,
330
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
Woosuk Kwon's avatar
Woosuk Kwon committed
331
    ) -> None:
332
        model_name = maybe_model_redirect(model_name)
333
        self.model_name = model_name
334

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

        model_kwargs = model_kwargs if model_kwargs is not None else {}
        model_kwargs.setdefault("torch_dtype", torch_dtype)

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

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

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

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

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

            self.model = model
389

390
391
392
393
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
394
                trust_remote_code=trust_remote_code,
395
            )
396

397
398
399
400
401
402
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
403
            trust_remote_code=trust_remote_code,
404
        )
405
406
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
407

408
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
409
        self,
410
        prompts: list[str],
411
        images: Optional[PromptImageInput] = None,
412
413
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
414
    ) -> list[Union[BatchFeature, BatchEncoding]]:
415
        if images is not None:
416
            assert len(prompts) == len(images)
417

418
419
420
421
422
423
        if videos is not None:
            assert len(prompts) == len(videos)

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

424
        all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
425
        for i, prompt in enumerate(prompts):
426
            processor_kwargs: dict[str, Any] = {
427
428
429
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
430
431
432
433
            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
434
435
436
437
438
439
440
441
442
            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
443
444

            inputs = self.processor(**processor_kwargs)
445
446
            if isinstance(inputs, BatchFeature):
                inputs = inputs.to(dtype=self.dtype)
447

448
449
450
451
            all_inputs.append(inputs)

        return all_inputs

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

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

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

        return outputs

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

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

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

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

    def generate_beam_search(
        self,
530
        prompts: list[str],
531
532
        beam_width: int,
        max_tokens: int,
533
534
535
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
536
    ) -> list[tuple[list[list[int]], list[str]]]:
537
538
539
540
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
541
542
543
544
545
                                num_return_sequences=beam_width,
                                images=images,
                                videos=videos,
                                audios=audios)

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

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

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

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

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

607
608
609
610
        return seq_logprobs

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

617
        # convert to dict
618
        seq_logprobs_lst: list[dict[int, float]] = []
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
        for tok_idx, tok_logprobs in enumerate(seq_logprobs):
            # drop prompt logprobs
            if tok_idx == 0:
                tok_logprobs = tok_logprobs[-1, :].reshape(1, -1)
            topk = tok_logprobs.topk(num_logprobs)

            tok_logprobs_dct = {}
            for token_id, logprob in zip(topk.indices[0], topk.values[0]):
                tok_logprobs_dct[token_id.item()] = logprob.item()

            seq_logprobs_lst.append(tok_logprobs_dct)

        return (
            seq_logprobs_lst,
            output_len,
        )

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

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

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

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

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

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

    def generate_encoder_decoder_greedy_logprobs_limit(
        self,
685
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
686
        max_tokens: int,
687
        num_logprobs: Optional[int],
688
        images: Optional[PromptImageInput] = None,
689
        **kwargs: Any,
690
    ) -> list[TokensTextLogprobs]:
691
692
693
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
694

695
696
697
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
698

699
700
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
701
            processor_kwargs: dict[str, Any] = {
702
703
704
705
706
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
707

708
709
            encoder_inputs = self.processor(**processor_kwargs)
            encoder_inputs = self.wrap_device(encoder_inputs)
710
711
712
713

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
714
715
716
                decoder_inputs = self.tokenizer(decoder_prompt,
                                                return_tensors="pt")
                decoder_input_ids = self.wrap_device(decoder_inputs.input_ids)
717
718
719
720
721
722
723
724

            output = self.model.generate(
                decoder_input_ids=decoder_input_ids,
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
725
                **encoder_inputs,
726
727
728
729
730
731
732
733
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
734
735
736
737
738
739
740
741
742
743
744

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))

        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

745
746
747
    def encode(self, prompts: list[str], *args,
               **kwargs) -> list[list[torch.Tensor]]:
        return self.model.encode(prompts, *args, **kwargs)
748

749
750
751
752
753
754
    def predict(self, prompts: list[list[str]], *args,
                **kwargs) -> torch.Tensor:
        return self.model.predict(prompts,
                                  *args,
                                  convert_to_tensor=True,
                                  **kwargs)
755

756
757
758
759
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
760
        del self.model
761
        cleanup_dist_env_and_memory()
762

Woosuk Kwon's avatar
Woosuk Kwon committed
763

Cyrus Leung's avatar
Cyrus Leung committed
764
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
765
766
767
768
769
def hf_runner():
    return HfRunner


class VllmRunner:
770
771
    """
    The default value of some arguments have been modified from
772
    {class}`~vllm.LLM` as follows:
773

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

    def __init__(
        self,
        model_name: str,
787
788
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
789
        tokenizer_name: Optional[str] = None,
790
        tokenizer_mode: str = "auto",
791
792
        trust_remote_code: bool = True,
        seed: Optional[int] = 0,
793
        max_model_len: Optional[int] = 1024,
794
        dtype: str = "auto",
795
        disable_log_stats: bool = True,
796
        tensor_parallel_size: int = 1,
797
        block_size: int = 16 if not torch.xpu.is_available() else 64,
798
        enable_chunked_prefill: Optional[bool] = False,
799
        swap_space: int = 4,
800
        enforce_eager: Optional[bool] = False,
801
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
802
    ) -> None:
803
        self.llm = LLM(
Woosuk Kwon's avatar
Woosuk Kwon committed
804
            model=model_name,
805
806
            runner=runner,
            convert=convert,
Woosuk Kwon's avatar
Woosuk Kwon committed
807
            tokenizer=tokenizer_name,
808
            tokenizer_mode=tokenizer_mode,
809
            trust_remote_code=trust_remote_code,
Woosuk Kwon's avatar
Woosuk Kwon committed
810
            dtype=dtype,
811
            seed=seed,
812
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
813
            enforce_eager=enforce_eager,
814
            disable_log_stats=disable_log_stats,
815
            tensor_parallel_size=tensor_parallel_size,
816
            max_model_len=max_model_len,
817
818
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
819
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
820
821
        )

822
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
823
        self,
824
        prompts: Union[list[str], list[torch.Tensor], list[int]],
825
        images: Optional[PromptImageInput] = None,
826
827
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
828
    ) -> list[TextPrompt]:
829

830
831
832
833
834
        if any(x is not None and len(x) != len(prompts)
               for x in [images, videos, audios]):
            raise ValueError(
                "All non-None multimodal inputs must have the same length as "
                "prompts")
835

836
837
838
839
840
841
842
843
844
845
        inputs = []
        for i, prompt in enumerate(prompts):
            multi_modal_data = {}
            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

846
            text_prompt_kwargs: dict[str, Any] = {
847
848
                "multi_modal_data": multi_modal_data or None
            }
849
850
851
852
853
854
855
            if isinstance(prompt, str):
                text_prompt_kwargs["prompt"] = prompt
            elif isinstance(prompt, list):
                text_prompt_kwargs["prompt_token_ids"] = prompt
            else:
                text_prompt_kwargs["prompt_embeds"] = prompt

856
            inputs.append(TextPrompt(**text_prompt_kwargs))
857
858
859
860
861

        return inputs

    def generate(
        self,
862
        prompts: Union[list[str], list[torch.Tensor]],
863
864
865
866
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
867
        **kwargs: Any,
868
    ) -> list[tuple[list[list[int]], list[str]]]:
869
870
871
872
873
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

874
875
876
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
877

878
        outputs: list[tuple[list[list[int]], list[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
879
880
881
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
882
883
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
884
885
            for sample in req_output.outputs:
                output_str = sample.text
886
                output_ids = list(sample.token_ids)
887
                req_sample_output_ids.append(prompt_ids + output_ids)
888
                req_sample_output_strs.append((prompt_str or "") + output_str)
889
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
890
891
        return outputs

892
    @staticmethod
893
    def _final_steps_generate_w_logprobs(
894
895
896
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
897
        for req_output in req_outputs:
898
            assert len(req_output.outputs) > 0
899
900
            for sample in req_output.outputs:
                output_str = sample.text
901
                output_ids = list(sample.token_ids)
902
                output_logprobs = sample.logprobs
903
904
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
905
906
        return outputs

907
908
    def generate_w_logprobs(
        self,
909
        prompts: list[str],
910
        sampling_params: SamplingParams,
911
912
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
913
        videos: Optional[PromptVideoInput] = None,
914
        **kwargs: Any,
915
916
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
917
918
919
920
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
921

922
923
924
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
925
926
927
928
929
930
931

        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        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)
932
933
934

    def generate_encoder_decoder_w_logprobs(
        self,
935
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
936
        sampling_params: SamplingParams,
937
938
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
939
940
941
942
943
        '''
        Logprobs generation for vLLM encoder/decoder models
        '''

        assert sampling_params.logprobs is not None
944
945
        req_outputs = self.llm.generate(encoder_decoder_prompts,
                                        sampling_params=sampling_params)
946
947
948
949
950
951
        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        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)
952

Woosuk Kwon's avatar
Woosuk Kwon committed
953
954
    def generate_greedy(
        self,
955
        prompts: Union[list[str], list[torch.Tensor]],
Woosuk Kwon's avatar
Woosuk Kwon committed
956
        max_tokens: int,
957
        images: Optional[PromptImageInput] = None,
958
959
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
960
        **kwargs: Any,
961
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
962
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
963
964
965
966
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
967
968
                                audios=audios,
                                **kwargs)
969
970
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
971

972
973
    def generate_greedy_logprobs(
        self,
974
        prompts: list[str],
975
        max_tokens: int,
976
        num_logprobs: Optional[int],
977
        num_prompt_logprobs: Optional[int] = None,
978
979
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
980
        videos: Optional[PromptVideoInput] = None,
981
982
        stop_token_ids: Optional[list[int]] = None,
        stop: Optional[list[str]] = None,
983
        **kwargs: Any,
984
985
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
986
987
988
989
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
990
            prompt_logprobs=num_prompt_logprobs,
991
992
            stop_token_ids=stop_token_ids,
            stop=stop)
993
994
995
996
997

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
998
999
                                        videos=videos,
                                        **kwargs)
1000

1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
    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
        """
        outputs = self.generate_greedy_logprobs(prompts,
                                                max_tokens=1,
                                                num_logprobs=None,
                                                num_prompt_logprobs=0)

        perplexities = []
        for output in outputs:
            output = cast(TokensTextLogprobsPromptLogprobs, output)
            token_datas = cast(list[Optional[dict[int, Logprob]]], output[3])
            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

1030
1031
    def generate_encoder_decoder_greedy_logprobs(
        self,
1032
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
1033
        max_tokens: int,
1034
        num_logprobs: Optional[int],
1035
        num_prompt_logprobs: Optional[int] = None,
1036
        skip_special_tokens: bool = True,
1037
1038
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
1039
1040
1041
1042
1043
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
1044
            skip_special_tokens=skip_special_tokens,
1045
        )
1046
1047
1048
1049
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

1050
        return self.generate_encoder_decoder_w_logprobs(
1051
1052
            encoder_decoder_prompts, greedy_logprobs_params)

1053
    def generate_beam_search(
1054
        self,
1055
        prompts: list[str],
1056
1057
        beam_width: int,
        max_tokens: int,
1058
1059
1060
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
1061
        concurrency_limit: Optional[int] = None,
1062
    ) -> list[tuple[list[list[int]], list[str]]]:
1063
1064
1065
1066
1067
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

1068
1069
1070
1071
        outputs = self.llm.beam_search(inputs,
                                       BeamSearchParams(beam_width=beam_width,
                                                        max_tokens=max_tokens),
                                       concurrency_limit=concurrency_limit)
1072
1073
1074
1075
1076
1077
1078
        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

1079
    def classify(self, prompts: list[str]) -> list[list[float]]:
1080
        req_outputs = self.llm.classify(prompts)
1081
1082
        return [req_output.outputs.probs for req_output in req_outputs]

1083
1084
1085
1086
1087
1088
1089
    def embed(self,
              prompts: list[str],
              images: Optional[PromptImageInput] = None,
              videos: Optional[PromptVideoInput] = None,
              audios: Optional[PromptAudioInput] = None,
              *args,
              **kwargs) -> list[list[float]]:
Cyrus Leung's avatar
Cyrus Leung committed
1090
1091
1092
1093
1094
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

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

1098
    def encode(self, prompts: list[str]) -> list[list[float]]:
1099
        req_outputs = self.llm.encode(prompts)
1100
1101
        return [req_output.outputs.data for req_output in req_outputs]

1102
1103
1104
1105
    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]

1106
1107
    def score(
        self,
1108
1109
        text_1: Union[str, list[str]],
        text_2: Union[str, list[str]],
1110
1111
        *args,
        **kwargs,
1112
    ) -> list[float]:
1113
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1114
        return [req_output.outputs.score for req_output in req_outputs]
1115

1116
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
        if hasattr(self.llm.llm_engine, "model_executor"):
            # This works either in V0 or in V1 with
            # VLLM_ENABLE_V1_MULTIPROCESSING=0
            executor = self.llm.llm_engine.model_executor
            return executor.apply_model(func)

        # This works in V1 with VLLM_ALLOW_INSECURE_SERIALIZATION=1
        def _apply_model(self):
            return func(self.get_model())

        return self.llm.llm_engine.collective_rpc(_apply_model)
1128

1129
1130
1131
    def get_llm(self) -> LLM:
        return self.llm

1132
1133
1134
1135
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1136
        del self.llm
1137
        cleanup_dist_env_and_memory()
1138

Woosuk Kwon's avatar
Woosuk Kwon committed
1139

1140
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1141
1142
def vllm_runner():
    return VllmRunner
1143
1144


1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
    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
1159
1160
1161
1162
1163
1164
1165


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

1166
1167
    from vllm.platforms import current_platform
    return current_platform.device_count()
1168
1169
1170


temp_dir = tempfile.gettempdir()
1171
1172
_dummy_opt_path = os.path.join(temp_dir, "dummy_opt")
_dummy_llava_path = os.path.join(temp_dir, "dummy_llava")
1173
_dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding")
1174
1175
1176
1177


@pytest.fixture
def dummy_opt_path():
1178
1179
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1180
        snapshot_download(repo_id="facebook/opt-125m",
1181
                          local_dir=_dummy_opt_path,
1182
1183
1184
1185
1186
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1187
        with open(json_path) as f:
1188
1189
1190
1191
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
    return _dummy_opt_path


@pytest.fixture
def dummy_llava_path():
    json_path = os.path.join(_dummy_llava_path, "config.json")
    if not os.path.exists(_dummy_llava_path):
        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"
                          ])
        assert os.path.exists(json_path)
1206
        with open(json_path) as f:
1207
1208
1209
1210
1211
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224


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


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
    parser.addoption("--optional",
                     action="store_true",
                     default=False,
                     help="run optional test")


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:
1249
            item.add_marker(skip_optional)
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261


@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")
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377


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
    server: Optional[http.server.ThreadingHTTPServer]
    thread: Optional[threading.Thread]

    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(
            (self.address, self.port), AssetHandler)
        self.thread = threading.Thread(target=self.server.serve_forever,
                                       daemon=True)
        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]:
    """
    Starts a thread based HTTP server bound to 127.0.0.1 on a random free port. 
    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]