utils.py 44.1 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
import contextlib
6
import copy
7
import functools
8
import importlib
9
import itertools
10
import json
11
import os
12
import random
13
import signal
14
15
import subprocess
import sys
16
import tempfile
17
import time
18
import warnings
19
from collections.abc import Callable, Iterable
20
from contextlib import ExitStack, contextmanager, suppress
21
from multiprocessing import Process
22
from pathlib import Path
23
from typing import Any, Literal
24
from unittest.mock import patch
25

26
import anthropic
27
import cloudpickle
28
import httpx
29
import openai
30
import pytest
31
import requests
32
import torch
33
import torch.nn.functional as F
34
from openai.types.completion import Completion
35
from typing_extensions import ParamSpec
36

37
import vllm.envs as envs
38
from tests.models.utils import TextTextLogprobs
39
40
41
42
from vllm.distributed import (
    ensure_model_parallel_initialized,
    init_distributed_environment,
)
43
from vllm.engine.arg_utils import AsyncEngineArgs
44
from vllm.entrypoints.cli.serve import ServeSubcommand
45
from vllm.model_executor.model_loader import get_model_loader
46
from vllm.platforms import current_platform
47
from vllm.transformers_utils.tokenizer import get_tokenizer
Cyrus Leung's avatar
Cyrus Leung committed
48
from vllm.utils.argparse_utils import FlexibleArgumentParser
49
from vllm.utils.mem_constants import GB_bytes
50
from vllm.utils.network_utils import get_open_port
51
from vllm.utils.torch_utils import cuda_device_count_stateless
52

53
if current_platform.is_rocm():
54
55
56
57
58
59
    from amdsmi import (
        amdsmi_get_gpu_vram_usage,
        amdsmi_get_processor_handles,
        amdsmi_init,
        amdsmi_shut_down,
    )
60
61
62
63
64
65
66
67

    @contextmanager
    def _nvml():
        try:
            amdsmi_init()
            yield
        finally:
            amdsmi_shut_down()
68
elif current_platform.is_cuda():
69
70
71
72
73
74
    from vllm.third_party.pynvml import (
        nvmlDeviceGetHandleByIndex,
        nvmlDeviceGetMemoryInfo,
        nvmlInit,
        nvmlShutdown,
    )
75
76
77
78
79
80
81
82

    @contextmanager
    def _nvml():
        try:
            nvmlInit()
            yield
        finally:
            nvmlShutdown()
83
84
85
86
87
else:

    @contextmanager
    def _nvml():
        yield
88

89

90
91
VLLM_PATH = Path(__file__).parent.parent
"""Path to root of the vLLM repository."""
92
93


94
95
class RemoteOpenAIServer:
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
96

97
    def _start_server(
98
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
99
100
    ) -> None:
        """Subclasses override this method to customize server process launch"""
101
102
103
        env = os.environ.copy()
        # the current process might initialize cuda,
        # to be safe, we should use spawn method
104
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
105
106
        if env_dict is not None:
            env.update(env_dict)
107
108
        serve_cmd = ["vllm", "serve", model, *vllm_serve_args]
        print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}")
109
        self.proc: subprocess.Popen = subprocess.Popen(
110
            serve_cmd,
111
112
113
114
115
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )

116
117
118
119
120
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
        *,
121
122
        env_dict: dict[str, str] | None = None,
        seed: int | None = 0,
123
        auto_port: bool = True,
124
125
        max_wait_seconds: float | None = None,
        override_hf_configs: dict[str, Any] | None = None,
126
    ) -> None:
127
        if auto_port:
128
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
129
130
131
                raise ValueError(
                    "You have manually specified the port when `auto_port=True`."
                )
132

133
134
135
            # No need for a port if using unix sockets
            if "--uds" not in vllm_serve_args:
                # Don't mutate the input args
136
                vllm_serve_args = vllm_serve_args + ["--port", str(get_open_port())]
137
138
        if seed is not None:
            if "--seed" in vllm_serve_args:
139
140
141
                raise ValueError(
                    f"You have manually specified the seed when `seed={seed}`."
                )
142
143

            vllm_serve_args = vllm_serve_args + ["--seed", str(seed)]
144

145
146
147
        if override_hf_configs is not None:
            vllm_serve_args = vllm_serve_args + [
                "--hf-overrides",
148
                json.dumps(override_hf_configs),
149
150
            ]

151
        parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")
152
153
        subparsers = parser.add_subparsers(required=False, dest="subparser")
        parser = ServeSubcommand().subparser_init(subparsers)
154
        args = parser.parse_args(["--model", model, *vllm_serve_args])
155
156
157
158
159
        self.uds = args.uds
        if args.uds:
            self.host = None
            self.port = None
        else:
160
            self.host = str(args.host or "127.0.0.1")
161
            self.port = int(args.port)
162

163
        self.show_hidden_metrics = args.show_hidden_metrics_for_version is not None
164

165
166
167
168
        # download the model before starting the server to avoid timeout
        is_local = os.path.isdir(model)
        if not is_local:
            engine_args = AsyncEngineArgs.from_cli_args(args)
169
170
171
172
173
            model_config = engine_args.create_model_config()
            load_config = engine_args.create_load_config()

            model_loader = get_model_loader(load_config)
            model_loader.download_model(model_config)
174

175
        self._start_server(model, vllm_serve_args, env_dict)
176
        max_wait_seconds = max_wait_seconds or 240
177
        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
178
179
180
181
182
183

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
184
        try:
185
            self.proc.wait(8)
186
187
188
        except subprocess.TimeoutExpired:
            # force kill if needed
            self.proc.kill()
189

190
    def _poll(self) -> int | None:
191
192
193
        """Subclasses override this method to customize process polling"""
        return self.proc.poll()

194
195
196
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
197
198
199
200
201
        client = (
            httpx.Client(transport=httpx.HTTPTransport(uds=self.uds))
            if self.uds
            else requests
        )
202
203
        while True:
            try:
204
                if client.get(url).status_code == 200:
205
                    break
206
207
208
209
210
            except Exception:
                # this exception can only be raised by requests.get,
                # which means the server is not ready yet.
                # the stack trace is not useful, so we suppress it
                # by using `raise from None`.
211
                result = self._poll()
212
                if result is not None and result != 0:
213
                    raise RuntimeError("Server exited unexpectedly.") from None
214
215
216

                time.sleep(0.5)
                if time.time() - start > timeout:
217
                    raise RuntimeError("Server failed to start in time.") from None
218
219
220

    @property
    def url_root(self) -> str:
221
222
223
224
225
        return (
            f"http://{self.uds.split('/')[-1]}"
            if self.uds
            else f"http://{self.host}:{self.port}"
        )
226
227
228
229

    def url_for(self, *parts: str) -> str:
        return self.url_root + "/" + "/".join(parts)

230
231
232
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
233
234
235
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
236
237
            max_retries=0,
            **kwargs,
238
239
        )

240
    def get_async_client(self, **kwargs):
241
242
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
243
244
245
246
247
248
        return openai.AsyncOpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )
249
250


251
252
253
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

254
    def _start_server(
255
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
256
    ) -> None:
257
        self.proc: Process = Process(
258
259
            target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
        )  # type: ignore[assignment]
260
261
        self.proc.start()

262
263
264
265
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
266
        child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
267
        *,
268
269
        env_dict: dict[str, str] | None = None,
        seed: int | None = 0,
270
        auto_port: bool = True,
271
        max_wait_seconds: float | None = None,
272
    ) -> None:
273
274
275
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
276
277
278
279
280
281
282
283
        super().__init__(
            model=model,
            vllm_serve_args=vllm_serve_args,
            env_dict=env_dict,
            seed=seed,
            auto_port=auto_port,
            max_wait_seconds=max_wait_seconds,
        )
284

285
    def _poll(self) -> int | None:
286
287
288
289
290
291
292
293
294
295
        return self.proc.exitcode

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
        self.proc.join(8)
        if self.proc.is_alive():
            # force kill if needed
            self.proc.kill()


296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
class RemoteAnthropicServer:
    DUMMY_API_KEY = "token-abc123"  # vLLM's Anthropic server does not need API key

    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
        *,
        env_dict: dict[str, str] | None = None,
        seed: int | None = 0,
        auto_port: bool = True,
        max_wait_seconds: float | None = None,
    ) -> None:
        if auto_port:
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
                raise ValueError(
                    "You have manually specified the port when `auto_port=True`."
                )

            # Don't mutate the input args
            vllm_serve_args = vllm_serve_args + ["--port", str(get_open_port())]
        if seed is not None:
            if "--seed" in vllm_serve_args:
                raise ValueError(
                    f"You have manually specified the seed when `seed={seed}`."
                )

            vllm_serve_args = vllm_serve_args + ["--seed", str(seed)]

        parser = FlexibleArgumentParser(description="vLLM's remote Anthropic server.")
        subparsers = parser.add_subparsers(required=False, dest="subparser")
        parser = ServeSubcommand().subparser_init(subparsers)
        args = parser.parse_args(["--model", model, *vllm_serve_args])
        self.host = str(args.host or "localhost")
        self.port = int(args.port)

        self.show_hidden_metrics = args.show_hidden_metrics_for_version is not None

        # download the model before starting the server to avoid timeout
        is_local = os.path.isdir(model)
        if not is_local:
            engine_args = AsyncEngineArgs.from_cli_args(args)
            model_config = engine_args.create_model_config()
            load_config = engine_args.create_load_config()

            model_loader = get_model_loader(load_config)
            model_loader.download_model(model_config)

        env = os.environ.copy()
        # the current process might initialize cuda,
        # to be safe, we should use spawn method
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
        if env_dict is not None:
            env.update(env_dict)
        self.proc = subprocess.Popen(
            [
                sys.executable,
                "-m",
                "vllm.entrypoints.anthropic.api_server",
                model,
                *vllm_serve_args,
            ],
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )
        max_wait_seconds = max_wait_seconds or 240
        self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
        try:
            self.proc.wait(8)
        except subprocess.TimeoutExpired:
            # force kill if needed
            self.proc.kill()

    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
        while True:
            try:
                if requests.get(url).status_code == 200:
                    break
            except Exception:
                # this exception can only be raised by requests.get,
                # which means the server is not ready yet.
                # the stack trace is not useful, so we suppress it
                # by using `raise from None`.
                result = self.proc.poll()
                if result is not None and result != 0:
                    raise RuntimeError("Server exited unexpectedly.") from None

                time.sleep(0.5)
                if time.time() - start > timeout:
                    raise RuntimeError("Server failed to start in time.") from None

    @property
    def url_root(self) -> str:
        return f"http://{self.host}:{self.port}"

    def url_for(self, *parts: str) -> str:
        return self.url_root + "/" + "/".join(parts)

    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
        return anthropic.Anthropic(
            base_url=self.url_for(),
            api_key=self.DUMMY_API_KEY,
            max_retries=0,
            **kwargs,
        )

    def get_async_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
        return anthropic.AsyncAnthropic(
            base_url=self.url_for(), api_key=self.DUMMY_API_KEY, max_retries=0, **kwargs
        )


421
422
423
424
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
425
    token_ids: list[int],
426
427
428
429
):
    results = []

    # test with text prompt
430
431
432
433
434
435
436
437
438
439
440
441
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=5, temperature=0.0
    )

    results.append(
        {
            "test": "single_completion",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
442
443
444
445
446
447
448
449
450

    # test using token IDs
    completion = client.completions.create(
        model=model,
        prompt=token_ids,
        max_tokens=5,
        temperature=0.0,
    )

451
452
453
454
455
456
457
458
    results.append(
        {
            "test": "token_ids",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
459
460

    # test seeded random sampling
461
462
463
464
465
466
467
468
469
470
471
472
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0
    )

    results.append(
        {
            "test": "seeded_sampling",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
473
474

    # test seeded random sampling with multiple prompts
475
476
477
478
479
480
481
482
483
484
485
486
    completion = client.completions.create(
        model=model, prompt=[prompt, prompt], max_tokens=5, seed=33, temperature=1.0
    )

    results.append(
        {
            "test": "seeded_sampling",
            "text": [choice.text for choice in completion.choices],
            "finish_reason": [choice.finish_reason for choice in completion.choices],
            "usage": completion.usage,
        }
    )
487
488
489
490
491
492
493
494
495

    # test simple list
    batch = client.completions.create(
        model=model,
        prompt=[prompt, prompt],
        max_tokens=5,
        temperature=0.0,
    )

496
497
498
499
500
501
502
    results.append(
        {
            "test": "simple_list",
            "text0": batch.choices[0].text,
            "text1": batch.choices[1].text,
        }
    )
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518

    # test streaming
    batch = client.completions.create(
        model=model,
        prompt=[prompt, prompt],
        max_tokens=5,
        temperature=0.0,
        stream=True,
    )

    texts = [""] * 2
    for chunk in batch:
        assert len(chunk.choices) == 1
        choice = chunk.choices[0]
        texts[choice.index] += choice.text

519
520
521
522
523
524
    results.append(
        {
            "test": "streaming",
            "texts": texts,
        }
    )
525
526
527
528

    return results


529
530
531
532
533
534
535
536
def _test_completion_close(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

    # test with text prompt
537
538
539
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=1, logprobs=5, temperature=0.0
    )
540

541
542
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
543

544
545
546
547
548
549
    results.append(
        {
            "test": "completion_close",
            "logprobs": logprobs,
        }
    )
550
551
552
553

    return results


554
555
556
557
558
559
560
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

561
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
562
563

    # test with text prompt
564
565
566
567
568
569
570
571
572
573
574
575
    chat_response = client.chat.completions.create(
        model=model, messages=messages, max_tokens=5, temperature=0.0
    )

    results.append(
        {
            "test": "completion_close",
            "text": chat_response.choices[0].message.content,
            "finish_reason": chat_response.choices[0].finish_reason,
            "usage": chat_response.usage,
        }
    )
576
577
578
579

    return results


580
581
582
583
584
585
586
587
588
589
590
591
592
593
def _test_embeddings(
    client: openai.OpenAI,
    model: str,
    text: str,
):
    results = []

    # test with text input
    embeddings = client.embeddings.create(
        model=model,
        input=text,
        encoding_format="float",
    )

594
595
596
597
598
599
600
    results.append(
        {
            "test": "single_embedding",
            "embedding": embeddings.data[0].embedding,
            "usage": embeddings.usage,
        }
    )
601
602
603
604

    return results


605
606
607
608
609
610
611
612
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "How do you feel today?"},
            ],
        }
    ]

    chat_completion = client.chat.completions.create(
        model=model_name,
        messages=messages,
        temperature=0.0,
        max_tokens=1,
        logprobs=True,
        top_logprobs=5,
    )
630
631
632
633
634
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

    for x in top_logprobs:
        x.logprob = round(x.logprob, 2)

635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
    results.append(
        {
            "test": "pure_text",
            "logprobs": top_logprobs,
        }
    )

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": image_url}},
                {"type": "text", "text": "What's in this image?"},
            ],
        }
    ]

    chat_completion = client.chat.completions.create(
        model=model_name,
        messages=messages,
        temperature=0.0,
        max_tokens=1,
        logprobs=True,
        top_logprobs=5,
    )
660
661
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

662
663
664
665
666
667
    results.append(
        {
            "test": "text_image",
            "logprobs": top_logprobs,
        }
    )
668
669
670
671

    return results


672
673
674
675
def compare_two_settings(
    model: str,
    arg1: list[str],
    arg2: list[str],
676
677
    env1: dict[str, str] | None = None,
    env2: dict[str, str] | None = None,
678
679
    *,
    method: str = "generate",
680
    max_wait_seconds: float | None = None,
681
) -> None:
682
    """
683
684
685
686
687
688
689
690
691
    Launch API server with two different sets of arguments/environments
    and compare the results of the API calls.

    Args:
        model: The model to test.
        arg1: The first set of arguments to pass to the API server.
        arg2: The second set of arguments to pass to the API server.
        env1: The first set of environment variables to pass to the API server.
        env2: The second set of environment variables to pass to the API server.
692
693
    """

694
695
696
697
698
699
700
701
702
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


703
704
705
def compare_all_settings(
    model: str,
    all_args: list[list[str]],
706
    all_envs: list[dict[str, str] | None],
707
708
    *,
    method: str = "generate",
709
    max_wait_seconds: float | None = None,
710
) -> None:
711
712
713
714
715
716
717
718
719
    """
    Launch API server with several different sets of arguments/environments
    and compare the results of the API calls with the first set of arguments.
    Args:
        model: The model to test.
        all_args: A list of argument lists to pass to the API server.
        all_envs: A list of environment dictionaries to pass to the API server.
    """

720
    trust_remote_code = False
721
    for args in all_args:
722
723
724
725
726
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
727
    for args in all_args:
728
729
730
731
732
733
734
735
736
        if "--tokenizer-mode" in args:
            tokenizer_mode = args[args.index("--tokenizer-mode") + 1]
            break

    tokenizer = get_tokenizer(
        model,
        trust_remote_code=trust_remote_code,
        tokenizer_mode=tokenizer_mode,
    )
737

738
739
740
741
742
743
744
    can_force_load_format = True

    for args in all_args:
        if "--load-format" in args:
            can_force_load_format = False
            break

745
    prompt = "Hello, my name is"
746
    token_ids = tokenizer(prompt).input_ids
747
    ref_results: list = []
748
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
749
750
751
752
753
754
755
756
757
        if can_force_load_format:
            # we are comparing the results and
            # usually we don't need real weights.
            # we force to use dummy weights by default,
            # and it should work for most of the cases.
            # if not, we can use VLLM_TEST_FORCE_LOAD_FORMAT
            # environment variable to force the load format,
            # e.g. in quantization tests.
            args = args + ["--load-format", envs.VLLM_TEST_FORCE_LOAD_FORMAT]
758
        compare_results: list = []
759
        results = ref_results if i == 0 else compare_results
760
761
762
        with RemoteOpenAIServer(
            model, args, env_dict=env, max_wait_seconds=max_wait_seconds
        ) as server:
763
764
765
766
767
768
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
769
770
771
772
773
774
775
            results.append(
                {
                    "test": "models_list",
                    "id": served_model.id,
                    "root": served_model.root,
                }
            )
776

777
778
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
779
780
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
781
782
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
783
784
            elif method == "generate_with_image":
                results += _test_image_text(
785
786
787
                    client,
                    model,
                    "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
788
                )
789
790
791
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
792
                raise ValueError(f"Unknown method: {method}")
793

794
795
796
797
798
799
            if i > 0:
                # if any setting fails, raise an error early
                ref_args = all_args[0]
                ref_envs = all_envs[0]
                compare_args = all_args[i]
                compare_envs = all_envs[i]
800
                for ref_result, compare_result in zip(ref_results, compare_results):
801
802
803
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
804
805
806
807
808
809
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
810
                            f"Embedding for {model=} are not the same.\n"
811
812
                            f"cosine_similarity={sim}\n"
                        )
813
814
                        del ref_result["embedding"]
                        del compare_result["embedding"]
815
816
817
818
819
                    assert ref_result == compare_result, (
                        f"Results for {model=} are not the same.\n"
                        f"{ref_args=} {ref_envs=}\n"
                        f"{compare_args=} {compare_envs=}\n"
                        f"{ref_result=}\n"
820
821
                        f"{compare_result=}\n"
                    )
822
823


824
825
826
827
828
829
830
831
832
833
834
835
def init_test_distributed_environment(
    tp_size: int,
    pp_size: int,
    rank: int,
    distributed_init_port: str,
    local_rank: int = -1,
) -> None:
    distributed_init_method = f"tcp://localhost:{distributed_init_port}"
    init_distributed_environment(
        world_size=pp_size * tp_size,
        rank=rank,
        distributed_init_method=distributed_init_method,
836
837
        local_rank=local_rank,
    )
838
839
840
    ensure_model_parallel_initialized(tp_size, pp_size)


841
def multi_process_parallel(
842
    monkeypatch: pytest.MonkeyPatch,
843
844
    tp_size: int,
    pp_size: int,
845
    test_target: Any,
846
) -> None:
847
848
    import ray

849
850
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
851
852
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
853
854
855
856
857
858
    # NOTE: Force ray not to use gitignore file as excluding, otherwise
    # it will not move .so files to working dir.
    # So we have to manually add some of large directories
    os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"
    ray.init(
        runtime_env={
859
            "working_dir": VLLM_PATH,
860
            "excludes": [
861
862
863
864
865
866
867
868
869
                "build",
                ".git",
                "cmake-build-*",
                "shellcheck",
                "dist",
                "ep_kernels_workspace",
            ],
        }
    )
870
871
872
873
874

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
875
876
877
878
879
880
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
881
882
            ),
        )
883
884
885
    ray.get(refs)

    ray.shutdown()
886
887
888


@contextmanager
889
def error_on_warning(category: type[Warning] = Warning):
890
891
    """
    Within the scope of this context manager, tests will fail if any warning
892
    of the given category is emitted.
893
894
    """
    with warnings.catch_warnings():
895
        warnings.filterwarnings("error", category=category)
896
897

        yield
898
899


900
901
902
903
904
905
906
907
908
909
def get_physical_device_indices(devices):
    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
    if visible_devices is None:
        return devices

    visible_indices = [int(x) for x in visible_devices.split(",")]
    index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
    return [index_mapping[i] for i in devices if i in index_mapping]


910
@_nvml()
911
912
913
def wait_for_gpu_memory_to_clear(
    *,
    devices: list[int],
914
915
    threshold_bytes: int | None = None,
    threshold_ratio: float | None = None,
916
917
    timeout_s: float = 120,
) -> None:
918
    assert threshold_bytes is not None or threshold_ratio is not None
919
920
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
921
    devices = get_physical_device_indices(devices)
922
923
    start_time = time.time()
    while True:
924
        output: dict[int, str] = {}
925
        output_raw: dict[int, tuple[float, float]] = {}
926
        for device in devices:
927
            if current_platform.is_rocm():
928
929
930
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
931
                gb_total = mem_info["vram_total"] / 2**10
932
933
934
935
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
936
937
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
938
            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
939

940
        print("gpu memory used/total (GiB): ", end="")
941
        for k, v in output.items():
942
943
            print(f"{k}={v}; ", end="")
        print("")
944

945
946
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
947
            threshold = f"{threshold_bytes / 2**30} GiB"
948
949
950
951
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

952
        dur_s = time.time() - start_time
953
        if all(is_free(used, total) for used, total in output_raw.values()):
954
955
956
957
            print(
                f"Done waiting for free GPU memory on devices {devices=} "
                f"({threshold=}) {dur_s=:.02f}"
            )
958
959
960
            break

        if dur_s >= timeout_s:
961
962
963
964
            raise ValueError(
                f"Memory of devices {devices=} not free after "
                f"{dur_s=:.02f} ({threshold=})"
            )
965
966

        time.sleep(5)
967
968


969
970
971
_P = ParamSpec("_P")


972
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
973
974
975
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
976

977
    @functools.wraps(func)
978
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
979
980
981
982
        # Make the process the leader of its own process group
        # to avoid sending SIGTERM to the parent process
        os.setpgrp()
        from _pytest.outcomes import Skipped
983
984
985

        # Create a unique temporary file to store exception info from child
        # process. Use test function name and process ID to avoid collisions.
986
987
        with (
            tempfile.NamedTemporaryFile(
988
                delete=False,
989
                mode="w+b",
990
                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
991
992
993
994
                suffix=".exc",
            ) as exc_file,
            ExitStack() as delete_after,
        ):
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
            exc_file_path = exc_file.name
            delete_after.callback(os.remove, exc_file_path)

            pid = os.fork()
            print(f"Fork a new process to run a test {pid}")
            if pid == 0:
                # Parent process responsible for deleting, don't delete
                # in child.
                delete_after.pop_all()
                try:
                    func(*args, **kwargs)
                except Skipped as e:
                    # convert Skipped to exit code 0
                    print(str(e))
                    os._exit(0)
                except Exception as e:
                    import traceback
1012

1013
1014
1015
1016
1017
1018
1019
                    tb_string = traceback.format_exc()

                    # Try to serialize the exception object first
                    exc_to_serialize: dict[str, Any]
                    try:
                        # First, try to pickle the actual exception with
                        # its traceback.
1020
                        exc_to_serialize = {"pickled_exception": e}
1021
1022
1023
1024
1025
                        # Test if it can be pickled
                        cloudpickle.dumps(exc_to_serialize)
                    except (Exception, KeyboardInterrupt):
                        # Fall back to string-based approach.
                        exc_to_serialize = {
1026
1027
1028
                            "exception_type": type(e).__name__,
                            "exception_msg": str(e),
                            "traceback": tb_string,
1029
1030
                        }
                    try:
1031
                        with open(exc_file_path, "wb") as f:
1032
1033
1034
1035
1036
1037
1038
                            cloudpickle.dump(exc_to_serialize, f)
                    except Exception:
                        # Fallback: just print the traceback.
                        print(tb_string)
                    os._exit(1)
                else:
                    os._exit(0)
1039
            else:
1040
1041
1042
                pgid = os.getpgid(pid)
                _pid, _exitcode = os.waitpid(pid, 0)
                # ignore SIGTERM signal itself
1043
                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
1044
1045
1046
1047
1048
1049
1050
1051
                # kill all child processes
                os.killpg(pgid, signal.SIGTERM)
                # restore the signal handler
                signal.signal(signal.SIGTERM, old_signal_handler)
                if _exitcode != 0:
                    # Try to read the exception from the child process
                    exc_info = {}
                    if os.path.exists(exc_file_path):
1052
1053
1054
1055
                        with (
                            contextlib.suppress(Exception),
                            open(exc_file_path, "rb") as f,
                        ):
1056
1057
                            exc_info = cloudpickle.load(f)

1058
1059
1060
                    if (
                        original_exception := exc_info.get("pickled_exception")
                    ) is not None:
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
                        # Re-raise the actual exception object if it was
                        # successfully pickled.
                        assert isinstance(original_exception, Exception)
                        raise original_exception

                    if (original_tb := exc_info.get("traceback")) is not None:
                        # Use string-based traceback for fallback case
                        raise AssertionError(
                            f"Test {func.__name__} failed when called with"
                            f" args {args} and kwargs {kwargs}"
                            f" (exit code: {_exitcode}):\n{original_tb}"
                        ) from None

                    # Fallback to the original generic error
                    raise AssertionError(
                        f"function {func.__name__} failed when called with"
                        f" args {args} and kwargs {kwargs}"
1078
1079
                        f" (exit code: {_exitcode})"
                    ) from None
1080
1081

    return wrapper
1082
1083


1084
1085
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function."""
1086
1087
1088
1089

    @functools.wraps(f)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
        # Check if we're already in a subprocess
1090
        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
1091
1092
1093
1094
            # If we are, just run the function directly
            return f(*args, **kwargs)

        import torch.multiprocessing as mp
1095

1096
        with suppress(RuntimeError):
1097
            mp.set_start_method("spawn")
1098
1099
1100
1101
1102
1103

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
1104
        env["RUNNING_IN_SUBPROCESS"] = "1"
1105
1106
1107
1108
1109
1110
1111

        with tempfile.TemporaryDirectory() as tempdir:
            output_filepath = os.path.join(tempdir, "new_process.tmp")

            # `cloudpickle` allows pickling complex functions directly
            input_bytes = cloudpickle.dumps((f, output_filepath))

1112
1113
1114
1115
1116
            repo_root = str(VLLM_PATH.resolve())

            env = dict(env or os.environ)
            env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")

1117
1118
            cmd = [sys.executable, "-m", f"{module_name}"]

1119
1120
1121
            returned = subprocess.run(
                cmd, input=input_bytes, capture_output=True, env=env
            )
1122
1123
1124
1125
1126
1127

            # check if the subprocess is successful
            try:
                returned.check_returncode()
            except Exception as e:
                # wrap raised exception to provide more information
1128
1129
1130
                raise RuntimeError(
                    f"Error raised in subprocess:\n{returned.stderr.decode()}"
                ) from e
1131
1132
1133
1134
1135

    return wrapper


def create_new_process_for_each_test(
1136
    method: Literal["spawn", "fork"] | None = None,
1137
1138
1139
1140
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
    """Creates a decorator that runs each test function in a new process.

    Args:
1141
        method: The process creation method. Can be either "spawn" or "fork".
1142
1143
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
1144
1145
1146
1147
1148

    Returns:
        A decorator to run test functions in separate processes.
    """
    if method is None:
1149
1150
        use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
        method = "spawn" if use_spawn else "fork"
1151

1152
    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1153
1154
1155
1156
1157
1158
1159

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


1160
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1161
1162
1163
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
1164

1165
1166
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
1167
1168
    """
    try:
1169
        if current_platform.is_cpu():
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
            memory_gb = 0
        else:
            memory_gb = current_platform.get_device_total_memory() / GB_bytes
    except Exception as e:
        warnings.warn(
            f"An error occurred when finding the available memory: {e}",
            stacklevel=2,
        )
        memory_gb = 0

1180
    return pytest.mark.skipif(
1181
        memory_gb < min_gb,
1182
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1183
1184
    )

1185
1186
1187
1188
1189
1190
1191
1192

def large_gpu_test(*, min_gb: int):
    """
    Decorate a test to be skipped if no GPU is available or it does not have
    sufficient memory.

    Currently, the CI machine uses L4 GPU which has 24 GB VRAM.
    """
1193
    mark = large_gpu_mark(min_gb)
1194

1195
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1196
        return mark(f)
1197
1198
1199
1200

    return wrapper


1201
1202
1203
def multi_gpu_marks(*, num_gpus: int):
    """Get a collection of pytest marks to apply for `@multi_gpu_test`."""
    test_selector = pytest.mark.distributed(num_gpus=num_gpus)
1204
1205
1206
1207
1208
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

1209
1210
1211
1212
1213
1214
1215
1216
1217
    return [test_selector, test_skipif]


def multi_gpu_test(*, num_gpus: int):
    """
    Decorate a test to be run only when multiple GPUs are available.
    """
    marks = multi_gpu_marks(num_gpus=num_gpus)

1218
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1219
        func = create_new_process_for_each_test()(f)
1220
1221
1222
1223
        for mark in reversed(marks):
            func = mark(func)

        return func
1224
1225
1226
1227

    return wrapper


1228
async def completions_with_server_args(
1229
    prompts: list[str],
1230
    model_name: str,
1231
    server_cli_args: list[str],
1232
    num_logprobs: int | None,
1233
    max_wait_seconds: int = 240,
1234
    max_tokens: int | list = 5,
1235
) -> list[Completion]:
1236
    """Construct a remote OpenAI server, obtain an async client to the
1237
1238
1239
1240
1241
1242
1243
1244
1245
    server & invoke the completions API to obtain completions.

    Args:
      prompts: test prompts
      model_name: model to spin up on the vLLM server
      server_cli_args: CLI args for starting the server
      num_logprobs: Number of logprobs to report (or `None`)
      max_wait_seconds: timeout interval for bringing up server.
                        Default: 240sec
1246
1247
1248
      max_tokens: max_tokens value for each of the given input prompts.
        if only one max_token value is given, the same value is used
        for all the prompts.
1249
1250
1251

    Returns:
      OpenAI Completion instance
1252
    """
1253

1254
1255
1256
1257
1258
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1259
    outputs = None
1260
1261
1262
    with RemoteOpenAIServer(
        model_name, server_cli_args, max_wait_seconds=max_wait_seconds
    ) as server:
1263
        client = server.get_async_client()
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
        outputs = [
            client.completions.create(
                model=model_name,
                prompt=[p],
                temperature=0,
                stream=False,
                max_tokens=max_tok,
                logprobs=num_logprobs,
            )
            for p, max_tok in zip(prompts, max_tokens)
        ]
1275
1276
        outputs = await asyncio.gather(*outputs)

1277
    assert outputs is not None, "Completion API call failed."
1278
1279
1280
1281

    return outputs


1282
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1283
    """Extract generated tokens from the output of a
1284
    request made to an Open-AI-protocol completions endpoint.
1285
    """
1286
1287
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1288
1289
1290


def get_client_text_logprob_generations(
1291
1292
1293
    completions: list[Completion],
) -> list[TextTextLogprobs]:
    """Operates on the output of a request made to an Open-AI-protocol
1294
    completions endpoint; obtains top-rank logprobs for each token in
1295
    each {class}`SequenceGroup`
1296
    """
1297
    text_generations = get_client_text_generations(completions)
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
    text = "".join(text_generations)
    return [
        (
            text_generations,
            text,
            (None if x.logprobs is None else x.logprobs.top_logprobs),
        )
        for completion in completions
        for x in completion.choices
    ]
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318


def has_module_attribute(module_name, attribute_name):
    """
    Helper function to check if a module has a specific attribute.
    """
    try:
        module = importlib.import_module(module_name)
        return hasattr(module, attribute_name)
    except ImportError:
        return False
1319
1320
1321
1322


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
1323
        return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]
1324
    elif current_platform.is_rocm():
1325
        attn_backend_list = ["TRITON_ATTN"]
1326
1327
        try:
            import aiter  # noqa: F401
1328

1329
            attn_backend_list.append("FLASH_ATTN")
1330
        except Exception:
1331
            print("Skip FLASH_ATTN on ROCm as aiter is not installed")
1332
1333

        return attn_backend_list
1334
1335
    elif current_platform.is_xpu():
        return ["FLASH_ATTN", "TRITON_ATTN"]
1336
1337
    else:
        raise ValueError("Unsupported platform")
1338
1339
1340
1341
1342


@contextmanager
def override_cutlass_fp8_supported(value: bool):
    with patch(
1343
1344
1345
        "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
        return_value=value,
    ):
1346
        yield
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365


def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
    """
    Generate prompts which a bunch of assignments,
    then asking for the value of one of them.
    The prompt is just under 10k tokens; sliding window is 4k
    so the answer is outside sliding window, but should still be correct.
    Args:
        batch_size: number of prompts to generate
        ln_range: an argument to control the length of the prompt
    """
    prompts: list[str] = []
    answer: list[int] = []
    indices: list[int] = []
    random.seed(1)
    for _ in range(batch_size):
        idx = random.randint(30, 90)
        indices.append(idx)
1366
1367
1368
1369
        prompt = (
            "```python\n# We set a number of variables, "
            + f"x{idx} will be important later\n"
        )
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
        ln = random.randint(*ln_range)
        for k in range(30, ln):
            v = random.randint(10, 99)
            if k == idx:
                answer.append(v)
            prompt += f"x{k} = {v}\n"
        prompt += f"# Now, we check the value of x{idx}:\n"
        prompt += f"assert x{idx} == "
        prompts.append(prompt)
    return prompts, answer, indices


1382
1383
1384
def check_answers(
    indices: list[int], answer: list[int], outputs: list[str], accept_rate: float = 0.7
):
1385
1386
1387
1388
1389
1390
1391
1392
1393
    answer2 = [int(text[0:2].strip()) for text in outputs]
    print(list(zip(indices, zip(answer, answer2))))
    numok = 0
    for a1, a2 in zip(answer, answer2):
        if a1 == a2:
            numok += 1
    frac_ok = numok / len(answer)
    print(f"Num OK: {numok}/{len(answer)} {frac_ok}")
    assert frac_ok >= accept_rate
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413


def flat_product(*iterables: Iterable[Any]):
    """
    Flatten lists of tuples of the cartesian product.
    Useful when we want to avoid nested tuples to allow
    test params to be unpacked directly from the decorator.

    Example:
    flat_product([(1, 2), (3, 4)], ["a", "b"]) ->
    [
      (1, 2, "a"),
      (1, 2, "b"),
      (3, 4, "a"),
      (3, 4, "b"),
    ]
    """
    for element in itertools.product(*iterables):
        normalized = (e if isinstance(e, tuple) else (e,) for e in element)
        yield tuple(itertools.chain(*normalized))