"tests/models/test_phi3v.py" did not exist on "fa9e3852290ecb6eaae45befbd629bb060f57fb7"
conftest.py 47.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12

# ruff: noqa

from tblib import pickling_support

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

13
import http.server
14
import json
15
import math
16
import mimetypes
17
import os
18
import socket
19
import tempfile
20
21
import threading
from collections.abc import Generator
22
from enum import Enum
23
from typing import Any, Callable, Optional, TypedDict, TypeVar, Union, cast
Woosuk Kwon's avatar
Woosuk Kwon committed
24

25
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
26
27
import pytest
import torch
28
import torch.nn as nn
29
import torch.nn.functional as F
30
from huggingface_hub import snapshot_download
31
from PIL import Image
32
33
from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
                          BatchEncoding, BatchFeature)
34
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
35

36
37
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
38
from vllm import LLM, SamplingParams
39
from vllm.assets.audio import AudioAsset
40
from vllm.assets.image import ImageAsset
41
from vllm.assets.video import VideoAsset
42
from vllm.config import ConvertOption, RunnerOption, _get_and_verify_dtype
43
from vllm.connections import global_http_connection
44
from vllm.distributed import (cleanup_dist_env_and_memory,
45
46
                              init_distributed_environment,
                              initialize_model_parallel)
47
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
48
                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
49
from vllm.logger import init_logger
50
from vllm.multimodal.utils import fetch_image
51
from vllm.outputs import RequestOutput
52
from vllm.sampling_params import BeamSearchParams
53
from vllm.sequence import Logprob
54
from vllm.transformers_utils.utils import maybe_model_redirect
55

56
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
57

58
59
60
_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")]
61
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
62

Cyrus Leung's avatar
Cyrus Leung committed
63
_M = TypeVar("_M")
64

65
_PromptMultiModalInput = Union[list[_M], list[list[_M]]]
Cyrus Leung's avatar
Cyrus Leung committed
66
67

PromptImageInput = _PromptMultiModalInput[Image.Image]
68
PromptAudioInput = _PromptMultiModalInput[tuple[np.ndarray, int]]
Cyrus Leung's avatar
Cyrus Leung committed
69
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
70

71

72
def _read_prompts(filename: str) -> list[str]:
73
    with open(filename) as f:
74
75
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
76
77


78
class ImageAssetPrompts(TypedDict):
79
80
    stop_sign: str
    cherry_blossom: str
81
82


83
class ImageTestAssets(list[ImageAsset]):
84
85

    def __init__(self) -> None:
86
87
88
89
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
90

91
    def prompts(self, prompts: ImageAssetPrompts) -> list[str]:
92
93
94
95
96
97
        """
        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.
        """
98
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
99
100


101
102
class VideoAssetPrompts(TypedDict):
    baby_reading: str
103
104


105
class VideoTestAssets(list[VideoAsset]):
106
107
108

    def __init__(self) -> None:
        super().__init__([
109
            VideoAsset("baby_reading"),
110
111
        ])

112
113
    def prompts(self, prompts: VideoAssetPrompts) -> list[str]:
        return [prompts["baby_reading"]]
114
115


116
class AudioAssetPrompts(TypedDict):
117
118
119
120
    mary_had_lamb: str
    winning_call: str


121
class AudioTestAssets(list[AudioAsset]):
122
123
124
125
126
127
128

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

129
    def prompts(self, prompts: AudioAssetPrompts) -> list[str]:
130
131
        return [prompts["mary_had_lamb"], prompts["winning_call"]]

132

133
IMAGE_ASSETS = ImageTestAssets()
134
"""Singleton instance of {class}`ImageTestAssets`."""
135
VIDEO_ASSETS = VideoTestAssets()
136
"""Singleton instance of {class}`VideoTestAssets`."""
137
AUDIO_ASSETS = AudioTestAssets()
138
"""Singleton instance of {class}`AudioTestAssets`."""
139
140


141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@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
161
@pytest.fixture(params=[True, False])
162
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
163
164
165
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
166
    skip_v0 = request.node.get_closest_marker("skip_v0")
Joe Runde's avatar
Joe Runde committed
167
168
169
170
171
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
172
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
173
    else:
174
175
        if skip_v0:
            pytest.skip("Skipping test on vllm V0")
176
177
178
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
179
180


181
182
183
184
185
186
187
@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


188
189
190
191
192
193
194
195
196
197
198
199
@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
200
    cleanup_dist_env_and_memory()
201
202


203
@pytest.fixture()
204
def should_do_global_cleanup_after_test(request) -> bool:
205
206
207
208
    """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.
    """
209

210
    return not request.node.get_closest_marker("skip_global_cleanup")
211
212


213
@pytest.fixture(autouse=True)
214
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
215
    yield
216
    if should_do_global_cleanup_after_test:
217
        cleanup_dist_env_and_memory()
218
219


220
221
222
223
224
225
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


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


234
235
236
237
238
239
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


240
241
242
243
244
245
246
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


247
@pytest.fixture
248
def example_encoder_decoder_prompts(
249
) -> dict[DecoderPromptType, list[ExplicitEncoderDecoderPrompt]]:
250
251
252
253
254
255
    '''
    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:
256

257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
    * 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:
272
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
273
        DecoderPromptType.EMPTY_STR:
274
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
275
        DecoderPromptType.CUSTOM:
276
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
277
278
279
    }


280
@pytest.fixture
281
def example_long_prompts() -> list[str]:
282
283
    prompts = []
    for filename in _LONG_PROMPTS:
284
        prompts += _read_prompts(filename)
285
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
286
287


288
@pytest.fixture(scope="session")
289
def image_assets() -> ImageTestAssets:
290
291
292
    return IMAGE_ASSETS


293
@pytest.fixture(scope="session")
294
def video_assets() -> VideoTestAssets:
295
296
297
    return VIDEO_ASSETS


298
@pytest.fixture(scope="session")
299
def audio_assets() -> AudioTestAssets:
300
301
302
    return AUDIO_ASSETS


303
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
304
_R = TypeVar("_R")
305

Woosuk Kwon's avatar
Woosuk Kwon committed
306
307
308

class HfRunner:

309
    def get_default_device(self):
310
        from vllm.platforms import current_platform
311

312
313
        return ("cpu"
                if current_platform.is_cpu() else current_platform.device_type)
314
315

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

319
        if device is None:
320
            device = self.device
321

322
323
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
324

325
326
327
328
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
329

Woosuk Kwon's avatar
Woosuk Kwon committed
330
331
332
    def __init__(
        self,
        model_name: str,
333
        dtype: str = "auto",
334
        *,
335
        model_kwargs: Optional[dict[str, Any]] = None,
336
        trust_remote_code: bool = True,
337
        is_sentence_transformer: bool = False,
338
        is_cross_encoder: bool = False,
339
        skip_tokenizer_init: bool = False,
340
        auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,
Woosuk Kwon's avatar
Woosuk Kwon committed
341
    ) -> None:
342
        model_name = maybe_model_redirect(model_name)
343
        self.model_name = model_name
344

345
346
        self.config = AutoConfig.from_pretrained(
            model_name,
347
            trust_remote_code=trust_remote_code,
348
349
        )
        self.device = self.get_default_device()
350
351
352
353
354
355
        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,
        )
356
357
358
359

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

360
        if is_sentence_transformer:
361
362
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
363
364
365
366
367

            self.model = SentenceTransformer(
                model_name,
                device=self.device,
                model_kwargs=model_kwargs,
368
                trust_remote_code=trust_remote_code,
369
            )
370
371
372
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
373
374
375
376
377

            self.model = CrossEncoder(
                model_name,
                device=self.device,
                automodel_args=model_kwargs,
378
                trust_remote_code=trust_remote_code,
379
            )
380
        else:
381
382
            model = auto_cls.from_pretrained(
                model_name,
383
                trust_remote_code=trust_remote_code,
384
385
386
                **model_kwargs,
            )

387
388
389
390
391
392
            # 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)

393
394
395
            if (getattr(model, "quantization_method", None) != "bitsandbytes"
                    and len({p.device
                             for p in model.parameters()}) < 2):
396
                model = model.to(device=self.device)
397
398

            self.model = model
399

400
401
402
403
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
404
                trust_remote_code=trust_remote_code,
405
            )
406

407
408
409
410
411
412
        # 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,
413
            trust_remote_code=trust_remote_code,
414
        )
415
416
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
417

418
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
419
        self,
420
        prompts: list[str],
421
        images: Optional[PromptImageInput] = None,
422
423
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
424
    ) -> list[Union[BatchFeature, BatchEncoding]]:
425
        if images is not None:
426
            assert len(prompts) == len(images)
427

428
429
430
431
432
433
        if videos is not None:
            assert len(prompts) == len(videos)

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

434
        all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
435
        for i, prompt in enumerate(prompts):
436
            processor_kwargs: dict[str, Any] = {
437
438
439
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
440
441
442
443
            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
444
445
446
447
448
449
450
451
452
            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
453
454

            inputs = self.processor(**processor_kwargs)
455
456
            if isinstance(inputs, BatchFeature):
                inputs = inputs.to(dtype=self.dtype)
457

458
459
460
461
            all_inputs.append(inputs)

        return all_inputs

462
463
464
465
466
467
468
469
470
    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

471
    def classify(self, prompts: list[str]) -> list[str]:
472
473
474
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
475
476
        problem_type = getattr(self.config, "problem_type", "")

477
478
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
479
480
481
482
483
484
            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()
485
486
487
488
            outputs.append(logits)

        return outputs

489
490
    def generate(
        self,
491
        prompts: list[str],
492
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
493
        videos: Optional[PromptVideoInput] = None,
494
495
        audios: Optional[PromptAudioInput] = None,
        **kwargs: Any,
496
    ) -> list[tuple[list[list[int]], list[str]]]:
497
498
499
500
501
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

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

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

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
537
538
539

    def generate_beam_search(
        self,
540
        prompts: list[str],
541
542
        beam_width: int,
        max_tokens: int,
543
544
545
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
546
    ) -> list[tuple[list[list[int]], list[str]]]:
547
548
549
550
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
551
552
553
554
555
                                num_return_sequences=beam_width,
                                images=images,
                                videos=videos,
                                audios=audios)

556
557
558
559
560
561
562
563
564
        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
565

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

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

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

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

617
618
619
620
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
621
        hidden_states: tuple[tuple[torch.Tensor, ...], ...],
622
        num_logprobs: Optional[int],
623
    ) -> tuple[list[dict[int, float]], int]:
624
625
626
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

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

646
647
    def generate_greedy_logprobs_limit(
        self,
648
        prompts: list[str],
649
        max_tokens: int,
650
        num_logprobs: Optional[int],
651
652
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
653
        videos: Optional[PromptVideoInput] = None,
654
        **kwargs: Any,
655
    ) -> list[TokensTextLogprobs]:
656
657
658
659
660
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

661
662
663
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
664

665
        for inputs in all_inputs:
666
            output = self.model.generate(
667
                **self.wrap_device(inputs),
668
669
670
671
672
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
673
                **kwargs,
674
675
            )

676
677
678
679
680
681
682
683
684
685
686
687
            (
                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))
688

689
690
691
692
693
694
        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,
695
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
696
        max_tokens: int,
697
        num_logprobs: Optional[int],
698
        images: Optional[PromptImageInput] = None,
699
        **kwargs: Any,
700
    ) -> list[TokensTextLogprobs]:
701
702
703
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
704

705
706
707
        all_logprobs: list[list[dict[int, float]]] = []
        all_output_ids: list[list[int]] = []
        all_output_strs: list[str] = []
708

709
710
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
711
            processor_kwargs: dict[str, Any] = {
712
713
714
715
716
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
717

718
719
            encoder_inputs = self.processor(**processor_kwargs)
            encoder_inputs = self.wrap_device(encoder_inputs)
720
721
722
723

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
724
725
726
                decoder_inputs = self.tokenizer(decoder_prompt,
                                                return_tensors="pt")
                decoder_input_ids = self.wrap_device(decoder_inputs.input_ids)
727
728
729
730
731
732
733
734

            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,
735
                **encoder_inputs,
736
737
738
739
740
741
742
743
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
744
745
746
747
748
749
750
751
752
753
754

            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]

755
756
757
    def encode(self, prompts: list[str], *args,
               **kwargs) -> list[list[torch.Tensor]]:
        return self.model.encode(prompts, *args, **kwargs)
758

759
760
761
762
763
764
    def predict(self, prompts: list[list[str]], *args,
                **kwargs) -> torch.Tensor:
        return self.model.predict(prompts,
                                  *args,
                                  convert_to_tensor=True,
                                  **kwargs)
765

766
767
768
769
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
770
        del self.model
771
        cleanup_dist_env_and_memory()
772

Woosuk Kwon's avatar
Woosuk Kwon committed
773

Cyrus Leung's avatar
Cyrus Leung committed
774
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
775
776
777
778
779
def hf_runner():
    return HfRunner


class VllmRunner:
780
781
    """
    The default value of some arguments have been modified from
782
    {class}`~vllm.LLM` as follows:
783

784
785
786
    - `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.
787
788
    - `block_size`: To reduce memory usage, set default to `64` if on XPU
        devices, otherwise default to `16`.
789
790
    - `enable_chunked_prefill`: Set to `False` instead of `None` for
      test reproducibility.
791
    - `enforce_eager`: Set to `False` to test CUDA graph.
792
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
793
794
795
796

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

832
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
833
        self,
834
        prompts: Union[list[str], list[torch.Tensor], list[int]],
835
        images: Optional[PromptImageInput] = None,
836
837
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
838
    ) -> list[TextPrompt]:
839

840
841
842
843
844
        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")
845

846
847
848
849
850
851
852
853
854
855
        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

856
            text_prompt_kwargs: dict[str, Any] = {
857
858
                "multi_modal_data": multi_modal_data or None
            }
859
860
861
862
863
864
865
            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

866
            inputs.append(TextPrompt(**text_prompt_kwargs))
867
868
869
870
871

        return inputs

    def generate(
        self,
872
        prompts: Union[list[str], list[torch.Tensor]],
873
874
875
876
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
877
        **kwargs: Any,
878
    ) -> list[tuple[list[list[int]], list[str]]]:
879
880
881
882
883
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

884
885
886
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
887

888
        outputs: list[tuple[list[list[int]], list[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
889
890
891
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
892
893
            req_sample_output_ids: list[list[int]] = []
            req_sample_output_strs: list[str] = []
894
895
            for sample in req_output.outputs:
                output_str = sample.text
896
                output_ids = list(sample.token_ids)
897
                req_sample_output_ids.append(prompt_ids + output_ids)
898
                req_sample_output_strs.append((prompt_str or "") + output_str)
899
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
900
901
        return outputs

902
    @staticmethod
903
    def _final_steps_generate_w_logprobs(
904
905
906
        req_outputs: list[RequestOutput],
    ) -> list[TokensTextLogprobsPromptLogprobs]:
        outputs: list[TokensTextLogprobsPromptLogprobs] = []
907
        for req_output in req_outputs:
908
            assert len(req_output.outputs) > 0
909
910
            for sample in req_output.outputs:
                output_str = sample.text
911
                output_ids = list(sample.token_ids)
912
                output_logprobs = sample.logprobs
913
914
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
915
916
        return outputs

917
918
    def generate_w_logprobs(
        self,
919
        prompts: list[str],
920
        sampling_params: SamplingParams,
921
922
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
923
        videos: Optional[PromptVideoInput] = None,
924
        **kwargs: Any,
925
926
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
927
928
929
930
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
931

932
933
934
        req_outputs = self.llm.generate(inputs,
                                        sampling_params=sampling_params,
                                        **kwargs)
935
936
937
938
939
940
941

        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)
942
943
944

    def generate_encoder_decoder_w_logprobs(
        self,
945
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
946
        sampling_params: SamplingParams,
947
948
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
949
950
951
952
953
        '''
        Logprobs generation for vLLM encoder/decoder models
        '''

        assert sampling_params.logprobs is not None
954
955
        req_outputs = self.llm.generate(encoder_decoder_prompts,
                                        sampling_params=sampling_params)
956
957
958
959
960
961
        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)
962

Woosuk Kwon's avatar
Woosuk Kwon committed
963
964
    def generate_greedy(
        self,
965
        prompts: Union[list[str], list[torch.Tensor]],
Woosuk Kwon's avatar
Woosuk Kwon committed
966
        max_tokens: int,
967
        images: Optional[PromptImageInput] = None,
968
969
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
970
        **kwargs: Any,
971
    ) -> list[tuple[list[int], str]]:
Woosuk Kwon's avatar
Woosuk Kwon committed
972
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
973
974
975
976
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
977
978
                                audios=audios,
                                **kwargs)
979
980
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
981

982
983
    def generate_greedy_logprobs(
        self,
984
        prompts: list[str],
985
        max_tokens: int,
986
        num_logprobs: Optional[int],
987
        num_prompt_logprobs: Optional[int] = None,
988
989
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
990
        videos: Optional[PromptVideoInput] = None,
991
992
        stop_token_ids: Optional[list[int]] = None,
        stop: Optional[list[str]] = None,
993
        **kwargs: Any,
994
995
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
996
997
998
999
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
1000
            prompt_logprobs=num_prompt_logprobs,
1001
1002
            stop_token_ids=stop_token_ids,
            stop=stop)
1003
1004
1005
1006
1007

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
1008
1009
                                        videos=videos,
                                        **kwargs)
1010

1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
    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

1040
1041
    def generate_encoder_decoder_greedy_logprobs(
        self,
1042
        encoder_decoder_prompts: list[ExplicitEncoderDecoderPrompt[str, str]],
1043
        max_tokens: int,
1044
        num_logprobs: Optional[int],
1045
        num_prompt_logprobs: Optional[int] = None,
1046
        skip_special_tokens: bool = True,
1047
1048
    ) -> Union[list[TokensTextLogprobs],
               list[TokensTextLogprobsPromptLogprobs]]:
1049
1050
1051
1052
1053
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
1054
            skip_special_tokens=skip_special_tokens,
1055
        )
1056
1057
1058
1059
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

1060
        return self.generate_encoder_decoder_w_logprobs(
1061
1062
            encoder_decoder_prompts, greedy_logprobs_params)

1063
    def generate_beam_search(
1064
        self,
1065
        prompts: list[str],
1066
1067
        beam_width: int,
        max_tokens: int,
1068
1069
1070
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
1071
        concurrency_limit: Optional[int] = None,
1072
    ) -> list[tuple[list[list[int]], list[str]]]:
1073
1074
1075
1076
1077
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

1078
1079
1080
1081
        outputs = self.llm.beam_search(inputs,
                                       BeamSearchParams(beam_width=beam_width,
                                                        max_tokens=max_tokens),
                                       concurrency_limit=concurrency_limit)
1082
1083
1084
1085
1086
1087
1088
        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

1089
    def classify(self, prompts: list[str]) -> list[list[float]]:
1090
        req_outputs = self.llm.classify(prompts)
1091
1092
        return [req_output.outputs.probs for req_output in req_outputs]

1093
1094
1095
1096
1097
1098
1099
    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
1100
1101
1102
1103
1104
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

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

1108
    def encode(self, prompts: list[str]) -> list[list[float]]:
1109
        req_outputs = self.llm.encode(prompts)
1110
1111
        return [req_output.outputs.data for req_output in req_outputs]

1112
1113
1114
1115
    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]

1116
1117
    def score(
        self,
1118
1119
        text_1: Union[str, list[str]],
        text_2: Union[str, list[str]],
1120
1121
        *args,
        **kwargs,
1122
    ) -> list[float]:
1123
        req_outputs = self.llm.score(text_1, text_2, *args, **kwargs)
1124
        return [req_output.outputs.score for req_output in req_outputs]
1125

1126
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
        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)
1138

1139
1140
1141
    def get_llm(self) -> LLM:
        return self.llm

1142
1143
1144
1145
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1146
        del self.llm
1147
        cleanup_dist_env_and_memory()
1148

Woosuk Kwon's avatar
Woosuk Kwon committed
1149

1150
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1151
1152
def vllm_runner():
    return VllmRunner
1153
1154


1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
@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
1169
1170
1171
1172
1173
1174
1175


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

1176
1177
    from vllm.platforms import current_platform
    return current_platform.device_count()
1178
1179
1180


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


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


@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)
1235
        with open(json_path) as f:
1236
1237
1238
1239
1240
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258


# 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:
1259
            item.add_marker(skip_optional)
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271


@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")
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
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387


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]