conftest.py 29.2 KB
Newer Older
1
2
import contextlib
import gc
3
import json
4
import os
5
import sys
6
import tempfile
7
from collections import UserList
8
from enum import Enum
9
10
from typing import (Any, Callable, Dict, List, Optional, Tuple, TypedDict,
                    TypeVar, Union)
Woosuk Kwon's avatar
Woosuk Kwon committed
11

12
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
13
14
import pytest
import torch
15
import torch.nn as nn
16
import torch.nn.functional as F
17
from huggingface_hub import snapshot_download
18
from PIL import Image
19
from transformers import (AutoModelForCausalLM, AutoTokenizer, BatchEncoding,
20
                          BatchFeature)
Woosuk Kwon's avatar
Woosuk Kwon committed
21
22

from vllm import LLM, SamplingParams
23
from vllm.assets.image import ImageAsset
24
from vllm.assets.video import VideoAsset
25
from vllm.config import TokenizerPoolConfig
26
from vllm.connections import global_http_connection
27
from vllm.distributed import (destroy_distributed_environment,
28
29
30
                              destroy_model_parallel,
                              init_distributed_environment,
                              initialize_model_parallel)
31
32
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
33
from vllm.logger import init_logger
34
from vllm.outputs import RequestOutput
35
from vllm.sequence import SampleLogprobs
36
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, cuda_device_count_stateless,
37
                        identity, is_cpu)
38

39
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
40

41
42
43
_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")]
44

45
46
47
PromptImageInput = Union[List[Image.Image], List[List[Image.Image]]]
PromptAudioInput = Union[List[Tuple[np.ndarray, int]],
                         List[List[Tuple[np.ndarray, int]]]]
48
PromptVideoInput = Union[List[np.ndarray], List[List[np.ndarray]]]
49

50

51
def _read_prompts(filename: str) -> List[str]:
52
    with open(filename, "r") as f:
53
54
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
55
56


57
58
59
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
60
61
62
63
64
65
66


if sys.version_info < (3, 9):
    # UserList cannot be subscripted
    class _ImageAssetsBase(UserList):
        pass
else:
67

68
69
    class _ImageAssetsBase(UserList[ImageAsset]):
        pass
70

71
72

class _ImageAssets(_ImageAssetsBase):
73
74

    def __init__(self) -> None:
75
76
77
78
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
79
80
81
82
83
84
85
86

    def prompts(self, prompts: _ImageAssetPrompts) -> List[str]:
        """
        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.
        """
87
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
88
89


90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


if sys.version_info < (3, 9):
    # UserList cannot be subscripted
    class _VideoAssetsBase(UserList):
        pass
else:

    class _VideoAssetsBase(UserList[VideoAsset]):
        pass


class _VideoAssets(_VideoAssetsBase):

    def __init__(self) -> None:
        super().__init__([
            VideoAsset("sample_demo_1.mp4"),
        ])

    def prompts(self, prompts: _VideoAssetPrompts) -> List[str]:
        return [prompts["sample_demo_1"]]


115
116
IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""
117
118
VIDEO_ASSETS = _VideoAssets()
"""Singleton instance of :class:`_VideoAssets`."""
119
120


121
122
123
124
125
126
127
@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


128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@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
    cleanup()


143
144
def cleanup():
    destroy_model_parallel()
145
    destroy_distributed_environment()
146
147
148
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
149
150
    if not is_cpu():
        torch.cuda.empty_cache()
151
152


153
@pytest.fixture()
154
def should_do_global_cleanup_after_test(request) -> bool:
155
156
157
158
    """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.
    """
159
160
161
162

    if request.node.get_closest_marker("skip_global_cleanup"):
        return False

163
164
165
    return True


166
@pytest.fixture(autouse=True)
167
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
168
    yield
169
170
    if should_do_global_cleanup_after_test:
        cleanup()
171
172


Woosuk Kwon's avatar
Woosuk Kwon committed
173
174
@pytest.fixture
def example_prompts() -> List[str]:
175
176
    prompts = []
    for filename in _TEST_PROMPTS:
177
        prompts += _read_prompts(filename)
178
179
180
    return prompts


181
182
183
184
185
186
187
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


188
@pytest.fixture
189
190
def example_encoder_decoder_prompts(
) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
191
192
193
194
195
196
    '''
    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:
197

198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
    * 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:
213
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
214
        DecoderPromptType.EMPTY_STR:
215
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
216
        DecoderPromptType.CUSTOM:
217
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
218
219
220
    }


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


229
230
231
232
233
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


234
235
236
237
238
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


239
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature)
240

Woosuk Kwon's avatar
Woosuk Kwon committed
241
242
243

class HfRunner:

244
    def wrap_device(self, input: _T) -> _T:
245
        if not is_cpu():
246
247
248
            # Check if the input is already on the GPU
            if hasattr(input, 'device') and input.device.type == "cuda":
                return input  # Already on GPU, no need to move
249
250
            return input.to("cuda")
        else:
251
252
253
            # Check if the input is already on the CPU
            if hasattr(input, 'device') and input.device.type == "cpu":
                return input  # Already on CPU, no need to move
254
255
            return input.to("cpu")

Woosuk Kwon's avatar
Woosuk Kwon committed
256
257
258
259
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
260
        *,
261
        model_kwargs: Optional[Dict[str, Any]] = None,
262
        is_embedding_model: bool = False,
263
        auto_cls=AutoModelForCausalLM,
264
265
        postprocess_inputs: Callable[[BatchEncoding],
                                     BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
266
    ) -> None:
267
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
268

269
        self.model_name = model_name
270

271
        if is_embedding_model:
272
273
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
274
275
276
277
278
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
                ).to(dtype=torch_dtype))
279
        else:
280
            model_kwargs = model_kwargs if model_kwargs is not None else {}
281
            self.model = self.wrap_device(
282
                auto_cls.from_pretrained(
283
284
285
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
286
                    **model_kwargs,
287
                ))
288
289
290
291
292
293
294
295

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
        )

        try:
296
297
298
            # don't put this import at the top level
            # it will call torch.cuda.device_count()
            from transformers import AutoProcessor  # noqa: F401
299
300
301
302
303
            self.processor = AutoProcessor.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
304
        except Exception as exc:
305
            logger.warning(
306
307
                "Unable to auto-load HuggingFace processor for model (%s). "
                "Using tokenizer instead. Reason: %s", model_name, exc)
308
            self.processor = self.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
309

310
311
        self.postprocess_inputs = postprocess_inputs

Woosuk Kwon's avatar
Woosuk Kwon committed
312
313
314
    def generate(
        self,
        prompts: List[str],
315
        images: Optional[PromptImageInput] = None,
316
        videos: Optional[List[np.ndarray]] = None,
317
        **kwargs: Any,
318
    ) -> List[Tuple[List[List[int]], List[str]]]:
319
320
        if images:
            assert len(prompts) == len(images)
321
322

        outputs: List[Tuple[List[List[int]], List[str]]] = []
323
        for i, prompt in enumerate(prompts):
324
325
326
327
328
329
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
330
331
            if videos is not None and videos[i] is not None:
                processor_kwargs["videos"] = videos[i]
332
333

            inputs = self.processor(**processor_kwargs)
334
            inputs = self.postprocess_inputs(inputs)
335

Woosuk Kwon's avatar
Woosuk Kwon committed
336
            output_ids = self.model.generate(
337
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
338
339
340
                use_cache=True,
                **kwargs,
            )
341
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
342
343
344
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
345
346
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
347
348
349
350
351
352
353
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
354
        images: Optional[PromptImageInput] = None,
355
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
356
    ) -> List[Tuple[List[int], str]]:
357
358
        outputs = self.generate(prompts,
                                do_sample=False,
359
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
360
361
                                images=images,
                                **kwargs)
362
363
364

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
365
366
367
368
369
370

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
371
    ) -> List[Tuple[List[List[int]], List[str]]]:
372
373
374
375
376
377
378
379
380
381
382
383
384
385
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
                                num_return_sequences=beam_width)
        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
386

387
388
389
390
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
391
        images: Optional[PromptImageInput] = None,
392
        videos: Optional[List[np.ndarray]] = None,
393
        **kwargs: Any,
394
    ) -> List[List[torch.Tensor]]:
395
396
397
398
399
400
401
402
        all_logprobs: List[List[torch.Tensor]] = []
        for i, prompt in enumerate(prompts):
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
403
404
            if videos is not None and videos[i] is not None:
                processor_kwargs["videos"] = videos[i]
405
406

            inputs = self.processor(**processor_kwargs)
407
            inputs = self.postprocess_inputs(inputs)
408

409
            output = self.model.generate(
410
                **self.wrap_device(inputs),
411
412
413
414
415
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
416
                **kwargs,
417
            )
418
            seq_logprobs: List[torch.Tensor] = []
419
420
421
422
423
424
425
426
427
            for hidden_states in output.hidden_states:
                last_hidden_states = hidden_states[-1][0]
                logits = torch.matmul(
                    last_hidden_states,
                    self.model.get_output_embeddings().weight.t(),
                )
                if self.model.get_output_embeddings().bias is not None:
                    logits += self.model.get_output_embeddings(
                    ).bias.unsqueeze(0)
428
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
429
430
431
432
                seq_logprobs.append(logprobs)
            all_logprobs.append(seq_logprobs)
        return all_logprobs

433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
    def _hidden_states_to_logprobs(
        self,
        hidden_states,
        num_logprobs,
    ) -> Tuple[List[Dict[int, float]], int]:
        seq_logprobs: List[torch.Tensor] = []
        output_len = len(hidden_states)
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
                last_hidden_states,
                self.model.get_output_embeddings().weight.t(),
            )
            if getattr(self.model.get_output_embeddings(), "bias",
                       None) is not None:
                logits += self.model.get_output_embeddings().bias.unsqueeze(0)
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

        # convert to dict
        seq_logprobs_lst: List[Dict[int, float]] = []
        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,
        )

471
472
473
474
475
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
476
477
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
478
        videos: Optional[List[np.ndarray]] = None,
479
        **kwargs: Any,
480
481
482
483
    ) -> List[Tuple[List[int], str, List[Dict[int, float]]]]:
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
484

485
486
487
488
489
490
491
492
        for i, prompt in enumerate(prompts):
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]

493
494
495
496
497
            if audios is not None:
                audio, sr = audios[i]
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr

498
499
            if videos is not None:
                processor_kwargs["videos"] = videos[i]
500
            inputs = self.processor(**processor_kwargs)
501
            inputs = self.postprocess_inputs(inputs)
502

503
            output = self.model.generate(
504
                **self.wrap_device(inputs),
505
506
507
508
509
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
510
                **kwargs,
511
512
            )

513
514
515
516
517
518
519
520
521
522
523
524
            (
                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))
525

526
527
528
529
530
531
        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,
532
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
533
534
535
536
537
538
539
        max_tokens: int,
        num_logprobs: int,
        **kwargs: Any,
    ) -> List[Tuple[List[int], str, List[Dict[int, float]]]]:
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
540

541
542
543
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
544

545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
        for (encoder_prompt,
             decoder_prompt) in to_enc_dec_tuple_list(encoder_decoder_prompts):
            encoder_input_ids = self.wrap_device(
                self.tokenizer(encoder_prompt, return_tensors="pt").input_ids)
            decoder_input_ids = (
                None if decoder_prompt is None else self.wrap_device(
                    self.tokenizer(decoder_prompt,
                                   return_tensors="pt").input_ids))

            output = self.model.generate(
                encoder_input_ids,
                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,
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
570
571
572
573
574
575
576
577
578
579
580

            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]

581
582
583
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

584
585
586
587
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
588
589
590
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
591

Cyrus Leung's avatar
Cyrus Leung committed
592
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
593
594
595
596
597
598
599
600
601
602
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
        tokenizer_name: Optional[str] = None,
603
604
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
605
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
606
        dtype: str = "half",
607
        disable_log_stats: bool = True,
608
        tensor_parallel_size: int = 1,
609
610
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
611
        swap_space: int = 4,
612
        enforce_eager: Optional[bool] = False,
613
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
614
615
616
617
618
619
    ) -> None:
        self.model = LLM(
            model=model_name,
            tokenizer=tokenizer_name,
            trust_remote_code=True,
            dtype=dtype,
620
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
621
            enforce_eager=enforce_eager,
622
            disable_log_stats=disable_log_stats,
623
            tensor_parallel_size=tensor_parallel_size,
624
            max_model_len=max_model_len,
625
626
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
627
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
628
629
630
631
632
633
        )

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
634
        images: Optional[PromptImageInput] = None,
635
    ) -> List[Tuple[List[List[int]], List[str]]]:
636
        if images is not None:
637
            assert len(prompts) == len(images)
638

639
640
641
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
642
                inputs[i]["multi_modal_data"] = {"image": image}
643

644
        req_outputs = self.model.generate(inputs,
645
                                          sampling_params=sampling_params)
646
647

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
648
649
650
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
651
652
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
653
654
            for sample in req_output.outputs:
                output_str = sample.text
655
                output_ids = list(sample.token_ids)
656
657
658
                req_sample_output_ids.append(prompt_ids + output_ids)
                req_sample_output_strs.append(prompt_str + output_str)
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
659
660
        return outputs

661
662
663
664
665
666
667
668
    def _final_steps_generate_w_logprobs(
        self,
        req_outputs: List[RequestOutput],
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
        outputs: List[Tuple[List[int], str, Optional[SampleLogprobs]]] = []
        for req_output in req_outputs:
            for sample in req_output.outputs:
                output_str = sample.text
669
                output_ids = list(sample.token_ids)
670
671
672
673
                output_logprobs = sample.logprobs
            outputs.append((output_ids, output_str, output_logprobs))
        return outputs

674
675
676
677
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
678
679
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
680
        videos: Optional[PromptVideoInput] = None,
681
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
682
683
        assert sampling_params.logprobs is not None

684
685
686
        if images is not None:
            assert len(prompts) == len(images)

687
688
689
        if videos is not None:
            assert len(prompts) == len(videos)

690
691
692
693
694
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
                inputs[i]["multi_modal_data"] = {"image": image}

695
696
697
698
        if audios is not None:
            for i, audio in enumerate(audios):
                inputs[i]["multi_modal_data"] = {"audio": audio}

699
700
701
702
703
        if videos is not None:
            for i, video in enumerate(videos):
                inputs[i]["multi_modal_data"] = {"video": video}
        print(f"[INPUTS!!!!]: {inputs}, {sampling_params}")

704
        req_outputs = self.model.generate(inputs,
705
                                          sampling_params=sampling_params)
706
707
708
709
        return self._final_steps_generate_w_logprobs(req_outputs)

    def generate_encoder_decoder_w_logprobs(
        self,
710
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
711
712
713
714
715
716
717
718
719
720
        sampling_params: SamplingParams,
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
        '''
        Logprobs generation for vLLM encoder/decoder models
        '''

        assert sampling_params.logprobs is not None
        req_outputs = self.model.generate(encoder_decoder_prompts,
                                          sampling_params=sampling_params)
        return self._final_steps_generate_w_logprobs(req_outputs)
721

Woosuk Kwon's avatar
Woosuk Kwon committed
722
723
724
725
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
726
        images: Optional[PromptImageInput] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
727
728
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
729
        outputs = self.generate(prompts, greedy_params, images=images)
730
731
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
732

733
734
735
736
737
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
738
739
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
740
        videos: Optional[PromptVideoInput] = None,
741
        stop_token_ids: Optional[List[int]] = None,
742
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
743
744
        greedy_logprobs_params = SamplingParams(temperature=0.0,
                                                max_tokens=max_tokens,
745
746
                                                logprobs=num_logprobs,
                                                stop_token_ids=stop_token_ids)
747
748
        outputs = self.generate_w_logprobs(prompts,
                                           greedy_logprobs_params,
749
                                           images=images,
750
751
                                           audios=audios,
                                           videos=videos)
752
753
754
755

        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

756
757
    def generate_encoder_decoder_greedy_logprobs(
        self,
758
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
        max_tokens: int,
        num_logprobs: int,
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
        greedy_logprobs_params = SamplingParams(temperature=0.0,
                                                use_beam_search=False,
                                                max_tokens=max_tokens,
                                                logprobs=num_logprobs)
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

        outputs = self.generate_encoder_decoder_w_logprobs(
            encoder_decoder_prompts, greedy_logprobs_params)

        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

776
777
778
779
780
    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
781
    ) -> List[Tuple[List[List[int]], List[str]]]:
782
783
784
785
786
787
        beam_search_params = SamplingParams(n=beam_width,
                                            use_beam_search=True,
                                            temperature=0.0,
                                            max_tokens=max_tokens)
        outputs = self.generate(prompts, beam_search_params)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
788

789
790
791
792
793
794
795
796
    def encode(self, prompts: List[str]) -> List[List[float]]:
        req_outputs = self.model.encode(prompts)
        outputs = []
        for req_output in req_outputs:
            embedding = req_output.outputs.embedding
            outputs.append(embedding)
        return outputs

797
798
799
800
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
801
802
803
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
804

805
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
806
807
def vllm_runner():
    return VllmRunner
808
809
810
811
812
813
814
815
816


def get_tokenizer_pool_config(tokenizer_group_type):
    if tokenizer_group_type is None:
        return None
    if tokenizer_group_type == "ray":
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type="ray",
                                   extra_config={})
817
818
819
820
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
821
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837


@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
838
839
840
841
842
843
844


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

845
    return cuda_device_count_stateless()
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868


temp_dir = tempfile.gettempdir()
_dummy_path = os.path.join(temp_dir, "dummy_opt")


@pytest.fixture
def dummy_opt_path():
    json_path = os.path.join(_dummy_path, "config.json")
    if not os.path.exists(_dummy_path):
        snapshot_download(repo_id="facebook/opt-125m",
                          local_dir=_dummy_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
        with open(json_path, "r") as f:
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_path