utils.py 40.5 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.tokenizers 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
        print(f"Environment variables: {env}")
110
        self.proc: subprocess.Popen = subprocess.Popen(
111
            serve_cmd,
112
113
114
115
116
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )

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

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

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

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

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

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

166
167
168
169
        # 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)
170
171
172
173
174
            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)
175

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

    def __enter__(self):
        return self

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

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

195
196
197
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
198
199
200
201
202
        client = (
            httpx.Client(transport=httpx.HTTPTransport(uds=self.uds))
            if self.uds
            else requests
        )
203
204
        while True:
            try:
205
                if client.get(url).status_code == 200:
206
                    break
207
208
209
210
211
            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`.
212
                result = self._poll()
213
                if result is not None and result != 0:
214
                    raise RuntimeError("Server exited unexpectedly.") from None
215
216
217

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

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

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

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

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

251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
    def get_client_anthropic(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_anthropic(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
        )

268

269
270
271
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

272
    def _start_server(
273
        self, model: str, vllm_serve_args: list[str], env_dict: dict[str, str] | None
274
    ) -> None:
275
        self.proc: Process = Process(
276
277
            target=self.child_process_fxn, args=(env_dict, model, vllm_serve_args)
        )  # type: ignore[assignment]
278
279
        self.proc.start()

280
281
282
283
    def __init__(
        self,
        model: str,
        vllm_serve_args: list[str],
284
        child_process_fxn: Callable[[dict[str, str] | None, str, list[str]], None],
285
        *,
286
        env_dict: dict[str, str] | None = None,
287
        seed: int = 0,
288
        auto_port: bool = True,
289
        max_wait_seconds: float | None = None,
290
    ) -> None:
291
292
293
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
294
295
296
297
298
299
300
301
        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,
        )
302

303
    def _poll(self) -> int | None:
304
305
306
307
308
309
310
311
312
313
        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()


314
315
316
317
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
318
    token_ids: list[int],
319
320
321
322
):
    results = []

    # test with text prompt
323
324
325
326
327
328
329
330
331
332
333
334
    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,
        }
    )
335
336
337
338
339
340
341
342
343

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

344
345
346
347
348
349
350
351
    results.append(
        {
            "test": "token_ids",
            "text": completion.choices[0].text,
            "finish_reason": completion.choices[0].finish_reason,
            "usage": completion.usage,
        }
    )
352
353

    # test seeded random sampling
354
355
356
357
358
359
360
361
362
363
364
365
    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,
        }
    )
366
367

    # test seeded random sampling with multiple prompts
368
369
370
371
372
373
374
375
376
377
378
379
    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,
        }
    )
380
381
382
383
384
385
386
387
388

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

389
390
391
392
393
394
395
    results.append(
        {
            "test": "simple_list",
            "text0": batch.choices[0].text,
            "text1": batch.choices[1].text,
        }
    )
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411

    # 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

412
413
414
415
416
417
    results.append(
        {
            "test": "streaming",
            "texts": texts,
        }
    )
418
419
420
421

    return results


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

    # test with text prompt
430
431
432
    completion = client.completions.create(
        model=model, prompt=prompt, max_tokens=1, logprobs=5, temperature=0.0
    )
433

434
435
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
436

437
438
439
440
441
442
    results.append(
        {
            "test": "completion_close",
            "logprobs": logprobs,
        }
    )
443
444
445
446

    return results


447
448
449
450
451
452
453
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

454
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
455
456

    # test with text prompt
457
458
459
460
461
462
463
464
465
466
467
468
    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,
        }
    )
469
470
471
472

    return results


473
474
475
476
477
478
479
480
481
482
483
484
485
486
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",
    )

487
488
489
490
491
492
493
    results.append(
        {
            "test": "single_embedding",
            "embedding": embeddings.data[0].embedding,
            "usage": embeddings.usage,
        }
    )
494
495
496
497

    return results


498
499
500
501
502
503
504
505
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
    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,
    )
523
524
525
526
527
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

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

528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
    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,
    )
553
554
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

555
556
557
558
559
560
    results.append(
        {
            "test": "text_image",
            "logprobs": top_logprobs,
        }
    )
561
562
563
564

    return results


565
566
567
568
def compare_two_settings(
    model: str,
    arg1: list[str],
    arg2: list[str],
569
570
    env1: dict[str, str] | None = None,
    env2: dict[str, str] | None = None,
571
572
    *,
    method: str = "generate",
573
    max_wait_seconds: float | None = None,
574
) -> None:
575
    """
576
577
578
579
580
581
582
583
584
    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.
585
586
    """

587
588
589
590
591
592
593
594
595
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


596
597
598
def compare_all_settings(
    model: str,
    all_args: list[list[str]],
599
    all_envs: list[dict[str, str] | None],
600
601
    *,
    method: str = "generate",
602
    max_wait_seconds: float | None = None,
603
) -> None:
604
605
606
607
608
609
610
611
612
    """
    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.
    """

613
    trust_remote_code = False
614
    for args in all_args:
615
616
617
618
619
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
620
    for args in all_args:
621
622
623
624
625
626
627
628
629
        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,
    )
630

631
632
633
634
635
636
637
    can_force_load_format = True

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

638
    prompt = "Hello, my name is"
639
    token_ids = tokenizer(prompt).input_ids
640
    ref_results: list = []
641
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
642
643
644
645
646
647
648
649
650
        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]
651
        compare_results: list = []
652
        results = ref_results if i == 0 else compare_results
653
654
655
        with RemoteOpenAIServer(
            model, args, env_dict=env, max_wait_seconds=max_wait_seconds
        ) as server:
656
657
658
659
660
661
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
662
663
664
665
666
667
668
            results.append(
                {
                    "test": "models_list",
                    "id": served_model.id,
                    "root": served_model.root,
                }
            )
669

670
671
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
672
673
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
674
675
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
676
677
            elif method == "generate_with_image":
                results += _test_image_text(
678
679
                    client,
                    model,
680
                    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
681
                )
682
683
684
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
685
                raise ValueError(f"Unknown method: {method}")
686

687
688
689
690
691
692
            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]
693
                for ref_result, compare_result in zip(ref_results, compare_results):
694
695
696
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
697
698
699
700
701
702
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
703
                            f"Embedding for {model=} are not the same.\n"
704
705
                            f"cosine_similarity={sim}\n"
                        )
706
707
                        del ref_result["embedding"]
                        del compare_result["embedding"]
708
709
710
711
712
                    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"
713
714
                        f"{compare_result=}\n"
                    )
715
716


717
718
719
720
721
722
723
724
725
726
727
728
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,
729
730
        local_rank=local_rank,
    )
731
732
733
    ensure_model_parallel_initialized(tp_size, pp_size)


734
def multi_process_parallel(
735
    monkeypatch: pytest.MonkeyPatch,
736
737
    tp_size: int,
    pp_size: int,
738
    test_target: Any,
739
) -> None:
740
741
    import ray

742
743
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
744
745
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
746
747
748
749
750
751
    # 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={
752
            "working_dir": VLLM_PATH,
753
            "excludes": [
754
755
756
757
758
759
760
761
762
                "build",
                ".git",
                "cmake-build-*",
                "shellcheck",
                "dist",
                "ep_kernels_workspace",
            ],
        }
    )
763
764
765
766
767

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
768
769
770
771
772
773
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
774
775
            ),
        )
776
777
778
    ray.get(refs)

    ray.shutdown()
779
780
781


@contextmanager
782
def error_on_warning(category: type[Warning] = Warning):
783
784
    """
    Within the scope of this context manager, tests will fail if any warning
785
    of the given category is emitted.
786
787
    """
    with warnings.catch_warnings():
788
        warnings.filterwarnings("error", category=category)
789
790

        yield
791
792


793
794
795
796
797
798
799
800
801
802
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]


803
@_nvml()
804
805
806
def wait_for_gpu_memory_to_clear(
    *,
    devices: list[int],
807
808
    threshold_bytes: int | None = None,
    threshold_ratio: float | None = None,
809
810
    timeout_s: float = 120,
) -> None:
811
    assert threshold_bytes is not None or threshold_ratio is not None
812
813
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
814
    devices = get_physical_device_indices(devices)
815
816
    start_time = time.time()
    while True:
817
        output: dict[int, str] = {}
818
        output_raw: dict[int, tuple[float, float]] = {}
819
        for device in devices:
820
            if current_platform.is_rocm():
821
822
823
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
824
                gb_total = mem_info["vram_total"] / 2**10
825
826
827
828
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
829
830
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
831
            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
832

833
        print("gpu memory used/total (GiB): ", end="")
834
        for k, v in output.items():
835
836
            print(f"{k}={v}; ", end="")
        print("")
837

838
839
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
840
            threshold = f"{threshold_bytes / 2**30} GiB"
841
842
843
844
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

845
        dur_s = time.time() - start_time
846
        if all(is_free(used, total) for used, total in output_raw.values()):
847
848
849
850
            print(
                f"Done waiting for free GPU memory on devices {devices=} "
                f"({threshold=}) {dur_s=:.02f}"
            )
851
852
853
            break

        if dur_s >= timeout_s:
854
855
856
857
            raise ValueError(
                f"Memory of devices {devices=} not free after "
                f"{dur_s=:.02f} ({threshold=})"
            )
858
859

        time.sleep(5)
860
861


862
863
864
_P = ParamSpec("_P")


865
def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
866
867
868
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
869

870
    @functools.wraps(func)
871
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
872
873
874
875
        # 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
876
877
878

        # Create a unique temporary file to store exception info from child
        # process. Use test function name and process ID to avoid collisions.
879
880
        with (
            tempfile.NamedTemporaryFile(
881
                delete=False,
882
                mode="w+b",
883
                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
884
885
886
887
                suffix=".exc",
            ) as exc_file,
            ExitStack() as delete_after,
        ):
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
            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
905

906
907
908
909
910
911
912
                    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.
913
                        exc_to_serialize = {"pickled_exception": e}
914
915
916
917
918
                        # Test if it can be pickled
                        cloudpickle.dumps(exc_to_serialize)
                    except (Exception, KeyboardInterrupt):
                        # Fall back to string-based approach.
                        exc_to_serialize = {
919
920
921
                            "exception_type": type(e).__name__,
                            "exception_msg": str(e),
                            "traceback": tb_string,
922
923
                        }
                    try:
924
                        with open(exc_file_path, "wb") as f:
925
926
927
928
929
930
931
                            cloudpickle.dump(exc_to_serialize, f)
                    except Exception:
                        # Fallback: just print the traceback.
                        print(tb_string)
                    os._exit(1)
                else:
                    os._exit(0)
932
            else:
933
934
935
                pgid = os.getpgid(pid)
                _pid, _exitcode = os.waitpid(pid, 0)
                # ignore SIGTERM signal itself
936
                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
937
938
939
940
941
942
943
944
                # 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):
945
946
947
948
                        with (
                            contextlib.suppress(Exception),
                            open(exc_file_path, "rb") as f,
                        ):
949
950
                            exc_info = cloudpickle.load(f)

951
952
953
                    if (
                        original_exception := exc_info.get("pickled_exception")
                    ) is not None:
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
                        # 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}"
971
972
                        f" (exit code: {_exitcode})"
                    ) from None
973
974

    return wrapper
975
976


977
978
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function."""
979
980
981
982

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

        import torch.multiprocessing as mp
988

989
        with suppress(RuntimeError):
990
            mp.set_start_method("spawn")
991
992
993
994
995
996

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
997
        env["RUNNING_IN_SUBPROCESS"] = "1"
998
999
1000
1001
1002
1003
1004

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

1005
1006
1007
1008
1009
            repo_root = str(VLLM_PATH.resolve())

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

1010
1011
            cmd = [sys.executable, "-m", f"{module_name}"]

1012
1013
1014
            returned = subprocess.run(
                cmd, input=input_bytes, capture_output=True, env=env
            )
1015
1016
1017
1018
1019
1020

            # check if the subprocess is successful
            try:
                returned.check_returncode()
            except Exception as e:
                # wrap raised exception to provide more information
1021
1022
1023
                raise RuntimeError(
                    f"Error raised in subprocess:\n{returned.stderr.decode()}"
                ) from e
1024
1025
1026
1027
1028

    return wrapper


def create_new_process_for_each_test(
1029
    method: Literal["spawn", "fork"] | None = None,
1030
1031
1032
1033
) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
    """Creates a decorator that runs each test function in a new process.

    Args:
1034
        method: The process creation method. Can be either "spawn" or "fork".
1035
1036
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
1037
1038
1039
1040
1041

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

1045
    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
1046
1047
1048
1049
1050
1051
1052

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


1053
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1054
1055
1056
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
1057

1058
1059
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
1060
1061
    """
    try:
1062
        if current_platform.is_cpu():
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
            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

1073
    return pytest.mark.skipif(
1074
        memory_gb < min_gb,
1075
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1076
1077
    )

1078

1079
1080
1081
1082
1083
1084
1085
requires_fp8 = pytest.mark.skipif(
    not current_platform.supports_fp8(),
    reason="FP8 is not supported on this GPU (requires Hopper or "
    "Ada architecture, compute capability 8.9+)",
)


1086
1087
1088
1089
1090
1091
1092
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.
    """
1093
    mark = large_gpu_mark(min_gb)
1094

1095
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1096
        return mark(f)
1097
1098
1099
1100

    return wrapper


1101
1102
1103
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)
1104
1105
1106
1107
1108
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

1109
1110
1111
1112
1113
1114
1115
1116
1117
    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)

1118
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
1119
        func = create_new_process_for_each_test()(f)
1120
1121
1122
1123
        for mark in reversed(marks):
            func = mark(func)

        return func
1124
1125
1126
1127

    return wrapper


1128
async def completions_with_server_args(
1129
    prompts: list[str],
1130
    model_name: str,
1131
    server_cli_args: list[str],
1132
    num_logprobs: int | None,
1133
    max_wait_seconds: int = 240,
1134
    max_tokens: int | list = 5,
1135
) -> list[Completion]:
1136
    """Construct a remote OpenAI server, obtain an async client to the
1137
1138
1139
1140
1141
1142
1143
1144
1145
    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
1146
1147
1148
      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.
1149
1150
1151

    Returns:
      OpenAI Completion instance
1152
    """
1153

1154
1155
1156
1157
1158
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1159
    outputs = None
1160
1161
1162
    with RemoteOpenAIServer(
        model_name, server_cli_args, max_wait_seconds=max_wait_seconds
    ) as server:
1163
        client = server.get_async_client()
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
        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)
        ]
1175
1176
        outputs = await asyncio.gather(*outputs)

1177
    assert outputs is not None, "Completion API call failed."
1178
1179
1180
1181

    return outputs


1182
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1183
    """Extract generated tokens from the output of a
1184
    request made to an Open-AI-protocol completions endpoint.
1185
    """
1186
1187
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1188
1189
1190


def get_client_text_logprob_generations(
1191
1192
1193
    completions: list[Completion],
) -> list[TextTextLogprobs]:
    """Operates on the output of a request made to an Open-AI-protocol
1194
    completions endpoint; obtains top-rank logprobs for each token in
1195
    each {class}`SequenceGroup`
1196
    """
1197
    text_generations = get_client_text_generations(completions)
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
    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
    ]
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218


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
1219
1220
1221
1222


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
1223
        return ["FLASH_ATTN", "TRITON_ATTN", "TREE_ATTN"]
1224
    elif current_platform.is_rocm():
1225
        attn_backend_list = ["TRITON_ATTN"]
1226
1227
        try:
            import aiter  # noqa: F401
1228

1229
            attn_backend_list.append("ROCM_AITER_FA")
1230
        except Exception:
1231
            print("Skip ROCM_AITER_FA on ROCm as aiter is not installed")
1232
1233

        return attn_backend_list
1234
1235
    elif current_platform.is_xpu():
        return ["FLASH_ATTN", "TRITON_ATTN"]
1236
1237
    else:
        raise ValueError("Unsupported platform")
1238
1239
1240
1241
1242


@contextmanager
def override_cutlass_fp8_supported(value: bool):
    with patch(
1243
1244
1245
        "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
        return_value=value,
    ):
1246
        yield
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265


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)
1266
1267
1268
1269
        prompt = (
            "```python\n# We set a number of variables, "
            + f"x{idx} will be important later\n"
        )
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
        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


1282
1283
1284
def check_answers(
    indices: list[int], answer: list[int], outputs: list[str], accept_rate: float = 0.7
):
1285
1286
1287
1288
1289
1290
1291
1292
1293
    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
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313


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