utils.py 35.4 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 copy
6
import functools
7
import importlib
8
import json
9
import os
10
import signal
11
12
import subprocess
import sys
13
import tempfile
14
import time
15
import warnings
16
from contextlib import contextmanager, suppress
17
from multiprocessing import Process
18
from pathlib import Path
19
from typing import Any, Callable, Literal, Optional, Union
20

21
import cloudpickle
22
import httpx
23
import openai
24
import pytest
25
import requests
26
import torch
27
import torch.nn.functional as F
28
from openai.types.completion import Completion
29
from typing_extensions import ParamSpec
30

31
import vllm.envs as envs
32
from tests.models.utils import TextTextLogprobs
33
34
from vllm.distributed import (ensure_model_parallel_initialized,
                              init_distributed_environment)
35
from vllm.engine.arg_utils import AsyncEngineArgs
36
from vllm.entrypoints.cli.serve import ServeSubcommand
37
from vllm.model_executor.model_loader import get_model_loader
38
from vllm.platforms import current_platform
39
from vllm.transformers_utils.tokenizer import get_tokenizer
40
from vllm.utils import (FlexibleArgumentParser, GB_bytes,
41
                        cuda_device_count_stateless, get_open_port)
42

43
if current_platform.is_rocm():
44
45
46
47
48
49
50
51
52
53
54
    from amdsmi import (amdsmi_get_gpu_vram_usage,
                        amdsmi_get_processor_handles, amdsmi_init,
                        amdsmi_shut_down)

    @contextmanager
    def _nvml():
        try:
            amdsmi_init()
            yield
        finally:
            amdsmi_shut_down()
55
elif current_platform.is_cuda():
56
57
58
    from vllm.third_party.pynvml import (nvmlDeviceGetHandleByIndex,
                                         nvmlDeviceGetMemoryInfo, nvmlInit,
                                         nvmlShutdown)
59
60
61
62
63
64
65
66

    @contextmanager
    def _nvml():
        try:
            nvmlInit()
            yield
        finally:
            nvmlShutdown()
67
68
69
70
71
else:

    @contextmanager
    def _nvml():
        yield
72

73

74
75
VLLM_PATH = Path(__file__).parent.parent
"""Path to root of the vLLM repository."""
76
77


78
79
class RemoteOpenAIServer:
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
80

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    def _start_server(self, model: str, vllm_serve_args: list[str],
                      env_dict: Optional[dict[str, str]]) -> None:
        """Subclasses override this method to customize server process launch
        """
        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 = subprocess.Popen(
            ["vllm", "serve", model, *vllm_serve_args],
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )

98
99
    def __init__(self,
                 model: str,
100
                 vllm_serve_args: list[str],
101
                 *,
102
                 env_dict: Optional[dict[str, str]] = None,
103
                 seed: Optional[int] = 0,
104
                 auto_port: bool = True,
105
106
                 max_wait_seconds: Optional[float] = None,
                 override_hf_configs: Optional[dict[str, Any]] = None) -> None:
107
        if auto_port:
108
109
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
                raise ValueError("You have manually specified the port "
110
111
                                 "when `auto_port=True`.")

112
113
114
115
116
117
            # No need for a port if using unix sockets
            if "--uds" not in vllm_serve_args:
                # Don't mutate the input args
                vllm_serve_args = vllm_serve_args + [
                    "--port", str(get_open_port())
                ]
118
119
120
121
122
123
        if seed is not None:
            if "--seed" in vllm_serve_args:
                raise ValueError("You have manually specified the seed "
                                 f"when `seed={seed}`.")

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

125
126
127
128
129
130
        if override_hf_configs is not None:
            vllm_serve_args = vllm_serve_args + [
                "--hf-overrides",
                json.dumps(override_hf_configs)
            ]

Ethan Xu's avatar
Ethan Xu committed
131
132
        parser = FlexibleArgumentParser(
            description="vLLM's remote OpenAI server.")
133
134
        subparsers = parser.add_subparsers(required=False, dest="subparser")
        parser = ServeSubcommand().subparser_init(subparsers)
135
        args = parser.parse_args(["--model", model, *vllm_serve_args])
136
137
138
139
140
141
142
        self.uds = args.uds
        if args.uds:
            self.host = None
            self.port = None
        else:
            self.host = str(args.host or 'localhost')
            self.port = int(args.port)
143

144
145
146
        self.show_hidden_metrics = \
            args.show_hidden_metrics_for_version is not None

147
148
149
150
        # 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)
151
152
153
154
155
            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)
156

157
        self._start_server(model, vllm_serve_args, env_dict)
158
        max_wait_seconds = max_wait_seconds or 240
159
        self._wait_for_server(url=self.url_for("health"),
160
                              timeout=max_wait_seconds)
161
162
163
164
165
166

    def __enter__(self):
        return self

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

173
174
175
176
    def _poll(self) -> Optional[int]:
        """Subclasses override this method to customize process polling"""
        return self.proc.poll()

177
178
179
    def _wait_for_server(self, *, url: str, timeout: float):
        # run health check
        start = time.time()
180
181
        client = (httpx.Client(transport=httpx.HTTPTransport(
            uds=self.uds)) if self.uds else requests)
182
183
        while True:
            try:
184
                if client.get(url).status_code == 200:
185
                    break
186
187
188
189
190
            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`.
191
                result = self._poll()
192
                if result is not None and result != 0:
193
                    raise RuntimeError("Server exited unexpectedly.") from None
194
195
196
197

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

    @property
    def url_root(self) -> str:
202
203
        return (f"http://{self.uds.split('/')[-1]}"
                if self.uds else f"http://{self.host}:{self.port}")
204
205
206
207

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

208
209
210
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
211
212
213
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
214
215
            max_retries=0,
            **kwargs,
216
217
        )

218
    def get_async_client(self, **kwargs):
219
220
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
221
222
223
224
        return openai.AsyncOpenAI(base_url=self.url_for("v1"),
                                  api_key=self.DUMMY_API_KEY,
                                  max_retries=0,
                                  **kwargs)
225
226


227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
class RemoteOpenAIServerCustom(RemoteOpenAIServer):
    """Launch test server with custom child process"""

    def _start_server(self, model: str, vllm_serve_args: list[str],
                      env_dict: Optional[dict[str, str]]) -> None:
        self.proc: Process = Process(
            target=self.child_process_fxn,
            args=(env_dict, model,
                  vllm_serve_args))  # type: ignore[assignment]
        self.proc.start()

    def __init__(self,
                 model: str,
                 vllm_serve_args: list[str],
                 child_process_fxn: Callable[
                     [Optional[dict[str, str]], str, list[str]], None],
                 *,
                 env_dict: Optional[dict[str, str]] = None,
                 seed: Optional[int] = 0,
                 auto_port: bool = True,
                 max_wait_seconds: Optional[float] = None) -> None:
        """Store custom child process function then invoke superclass
        constructor which will indirectly launch it."""
        self.child_process_fxn = child_process_fxn
        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)

    def _poll(self) -> Optional[int]:
        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()


269
270
271
272
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
273
    token_ids: list[int],
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
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
):
    results = []

    # test with text prompt
    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,
    })

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

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

    # test seeded random sampling
    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,
    })

    # test seeded random sampling with multiple prompts
    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,
    })

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

    results.append({
        "test": "simple_list",
        "text0": batch.choices[0].text,
        "text1": batch.choices[1].text,
    })

    # 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

    results.append({
        "test": "streaming",
        "texts": texts,
    })

    return results


373
374
375
376
377
378
379
380
381
382
383
384
385
386
def _test_completion_close(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

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

387
388
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
389
390
391

    results.append({
        "test": "completion_close",
392
        "logprobs": logprobs,
393
394
395
396
397
    })

    return results


398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
def _test_chat(
    client: openai.OpenAI,
    model: str,
    prompt: str,
):
    results = []

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

    # test with text prompt
    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,
    })

    return results


429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
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",
    )

    results.append({
        "test": "single_embedding",
        "embedding": embeddings.data[0].embedding,
        "usage": embeddings.usage,
    })

    return results


452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
def _test_image_text(
    client: openai.OpenAI,
    model_name: str,
    image_url: str,
):
    results = []

    # test pure text input
    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)
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

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

    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)
    top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs

    results.append({
        "test": "text_image",
        "logprobs": top_logprobs,
    })

    return results


520
def compare_two_settings(model: str,
521
522
523
524
                         arg1: list[str],
                         arg2: list[str],
                         env1: Optional[dict[str, str]] = None,
                         env2: Optional[dict[str, str]] = None,
525
                         *,
526
                         method: str = "generate",
527
                         max_wait_seconds: Optional[float] = None) -> None:
528
    """
529
530
531
532
533
534
535
536
537
    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.
538
539
    """

540
541
542
543
544
545
546
547
548
549
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


def compare_all_settings(model: str,
550
551
                         all_args: list[list[str]],
                         all_envs: list[Optional[dict[str, str]]],
552
                         *,
553
                         method: str = "generate",
554
555
556
557
558
559
560
561
562
563
                         max_wait_seconds: Optional[float] = None) -> None:
    """
    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.
    """

564
    trust_remote_code = False
565
    for args in all_args:
566
567
568
569
570
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
571
    for args in all_args:
572
573
574
575
576
577
578
579
580
        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,
    )
581

582
583
584
585
586
587
588
    can_force_load_format = True

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

589
    prompt = "Hello, my name is"
590
    token_ids = tokenizer(prompt).input_ids
591
    ref_results: list = []
592
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
593
594
595
596
597
598
599
600
601
        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]
602
        compare_results: list = []
603
        results = ref_results if i == 0 else compare_results
604
605
606
607
        with RemoteOpenAIServer(model,
                                args,
                                env_dict=env,
                                max_wait_seconds=max_wait_seconds) as server:
608
609
610
611
612
613
614
615
616
617
618
619
            client = server.get_client()

            # test models list
            models = client.models.list()
            models = models.data
            served_model = models[0]
            results.append({
                "test": "models_list",
                "id": served_model.id,
                "root": served_model.root,
            })

620
621
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
622
623
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
624
625
            elif method == "generate_chat":
                results += _test_chat(client, model, prompt)
626
627
628
629
630
            elif method == "generate_with_image":
                results += _test_image_text(
                    client, model,
                    "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png"
                )
631
632
633
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
634
                raise ValueError(f"Unknown method: {method}")
635

636
637
638
639
640
641
642
643
            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]
                for ref_result, compare_result in zip(ref_results,
                                                      compare_results):
644
645
646
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
647
648
649
650
651
652
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
653
                            f"Embedding for {model=} are not the same.\n"
654
                            f"cosine_similarity={sim}\n")
655
656
                        del ref_result["embedding"]
                        del compare_result["embedding"]
657
658
659
660
661
662
                    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"
                        f"{compare_result=}\n")
663
664


665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
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,
        local_rank=local_rank)
    ensure_model_parallel_initialized(tp_size, pp_size)


681
def multi_process_parallel(
682
    monkeypatch: pytest.MonkeyPatch,
683
684
    tp_size: int,
    pp_size: int,
685
    test_target: Any,
686
) -> None:
687
688
    import ray

689
690
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
691
692
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
693
694
695
696
697
698
699
700
701
702
    # 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={
            "working_dir": VLLM_PATH,
            "excludes":
            ["build", ".git", "cmake-build-*", "shellcheck", "dist"]
        })
703
704
705
706
707

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
708
709
710
711
712
713
714
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
            ), )
715
716
717
    ray.get(refs)

    ray.shutdown()
718
719
720


@contextmanager
721
def error_on_warning(category: type[Warning] = Warning):
722
723
    """
    Within the scope of this context manager, tests will fail if any warning
724
    of the given category is emitted.
725
726
    """
    with warnings.catch_warnings():
727
        warnings.filterwarnings("error", category=category)
728
729

        yield
730
731


732
733
734
735
736
737
738
739
740
741
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]


742
@_nvml()
743
744
745
746
def wait_for_gpu_memory_to_clear(*,
                                 devices: list[int],
                                 threshold_bytes: Optional[int] = None,
                                 threshold_ratio: Optional[float] = None,
747
                                 timeout_s: float = 120) -> None:
748
    assert threshold_bytes is not None or threshold_ratio is not None
749
750
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
751
    devices = get_physical_device_indices(devices)
752
753
    start_time = time.time()
    while True:
754
        output: dict[int, str] = {}
755
        output_raw: dict[int, tuple[float, float]] = {}
756
        for device in devices:
757
            if current_platform.is_rocm():
758
759
760
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
761
                gb_total = mem_info["vram_total"] / 2**10
762
763
764
765
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
766
767
768
                gb_total = mem_info.total / 2**30
            output_raw[device] = (gb_used, gb_total)
            output[device] = f'{gb_used:.02f}/{gb_total:.02f}'
769

770
        print('gpu memory used/total (GiB): ', end='')
771
772
773
774
        for k, v in output.items():
            print(f'{k}={v}; ', end='')
        print('')

775
776
777
778
779
780
781
        if threshold_bytes is not None:
            is_free = lambda used, total: used <= threshold_bytes / 2**30
            threshold = f"{threshold_bytes/2**30} GiB"
        else:
            is_free = lambda used, total: used / total <= threshold_ratio
            threshold = f"{threshold_ratio:.2f}"

782
        dur_s = time.time() - start_time
783
        if all(is_free(used, total) for used, total in output_raw.values()):
784
            print(f'Done waiting for free GPU memory on devices {devices=} '
785
                  f'({threshold=}) {dur_s=:.02f}')
786
787
788
789
            break

        if dur_s >= timeout_s:
            raise ValueError(f'Memory of devices {devices=} not free after '
790
                             f'{dur_s=:.02f} ({threshold=})')
791
792

        time.sleep(5)
793
794


795
796
797
798
799
_P = ParamSpec("_P")


def fork_new_process_for_each_test(
        f: Callable[_P, None]) -> Callable[_P, None]:
800
801
802
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
803
804

    @functools.wraps(f)
805
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
806
807
808
809
810
        # 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
        pid = os.fork()
811
        print(f"Fork a new process to run a test {pid}")
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
        if pid == 0:
            try:
                f(*args, **kwargs)
            except Skipped as e:
                # convert Skipped to exit code 0
                print(str(e))
                os._exit(0)
            except Exception:
                import traceback
                traceback.print_exc()
                os._exit(1)
            else:
                os._exit(0)
        else:
            pgid = os.getpgid(pid)
            _pid, _exitcode = os.waitpid(pid, 0)
            # ignore SIGTERM signal itself
829
            old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
830
831
832
            # kill all child processes
            os.killpg(pgid, signal.SIGTERM)
            # restore the signal handler
833
            signal.signal(signal.SIGTERM, old_signal_handler)
834
835
836
837
            assert _exitcode == 0, (f"function {f} failed when called with"
                                    f" args {args} and kwargs {kwargs}")

    return wrapper
838
839


840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
def spawn_new_process_for_each_test(
        f: Callable[_P, None]) -> Callable[_P, None]:
    """Decorator to spawn a new process for each test function.
    """

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

        import torch.multiprocessing as mp
        with suppress(RuntimeError):
            mp.set_start_method('spawn')

        # Get the module
        module_name = f.__module__

        # Create a process with environment variable set
        env = os.environ.copy()
        env['RUNNING_IN_SUBPROCESS'] = '1'

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

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

            returned = subprocess.run(cmd,
                                      input=input_bytes,
                                      capture_output=True,
                                      env=env)

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

    return wrapper


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

    Args:
        method: The process creation method. Can be either "spawn" or "fork". 
894
895
               If not specified, it defaults to "spawn" on ROCm and XPU
               platforms and "fork" otherwise.
896
897
898
899
900

    Returns:
        A decorator to run test functions in separate processes.
    """
    if method is None:
901
902
        use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
        method = "spawn" if use_spawn else "fork"
903
904
905
906
907
908
909
910
911
912

    assert method in ["spawn",
                      "fork"], "Method must be either 'spawn' or 'fork'"

    if method == "fork":
        return fork_new_process_for_each_test

    return spawn_new_process_for_each_test


913
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
914
915
916
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
917

918
919
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
920
921
    """
    try:
922
        if current_platform.is_cpu():
923
924
925
926
927
928
929
930
931
932
            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

933
    return pytest.mark.skipif(
934
        memory_gb < min_gb,
935
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
936
937
    )

938
939
940
941
942
943
944
945

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.
    """
946
    mark = large_gpu_mark(min_gb)
947

948
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
949
        return mark(f)
950
951
952
953

    return wrapper


954
955
956
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)
957
958
959
960
961
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

962
963
964
965
966
967
968
969
970
    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)

971
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
972
        func = create_new_process_for_each_test()(f)
973
974
975
976
        for mark in reversed(marks):
            func = mark(func)

        return func
977
978
979
980

    return wrapper


981
async def completions_with_server_args(
982
    prompts: list[str],
983
    model_name: str,
984
    server_cli_args: list[str],
985
986
    num_logprobs: Optional[int],
    max_wait_seconds: int = 240,
987
    max_tokens: Union[int, list] = 5,
988
) -> list[Completion]:
989
990
991
992
993
994
995
996
997
998
    '''Construct a remote OpenAI server, obtain an async client to the
    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
999
1000
1001
      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.
1002
1003
1004
1005
1006

    Returns:
      OpenAI Completion instance
    '''

1007
1008
1009
1010
1011
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

1012
1013
1014
1015
1016
    outputs = None
    with RemoteOpenAIServer(model_name,
                            server_cli_args,
                            max_wait_seconds=max_wait_seconds) as server:
        client = server.get_async_client()
1017
1018
1019
1020
1021
1022
1023
1024
1025
        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) ]
        outputs = await asyncio.gather(*outputs)

1026
    assert outputs is not None, "Completion API call failed."
1027
1028
1029
1030

    return outputs


1031
def get_client_text_generations(completions: list[Completion]) -> list[str]:
1032
1033
1034
    '''Extract generated tokens from the output of a
    request made to an Open-AI-protocol completions endpoint.
    '''
1035
1036
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
1037
1038
1039


def get_client_text_logprob_generations(
1040
        completions: list[Completion]) -> list[TextTextLogprobs]:
1041
1042
    '''Operates on the output of a request made to an Open-AI-protocol
    completions endpoint; obtains top-rank logprobs for each token in
1043
    each {class}`SequenceGroup`
1044
1045
1046
1047
1048
    '''
    text_generations = get_client_text_generations(completions)
    text = ''.join(text_generations)
    return [(text_generations, text,
             (None if x.logprobs is None else x.logprobs.top_logprobs))
1049
            for completion in completions for x in completion.choices]
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060


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
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076


def get_attn_backend_list_based_on_platform() -> list[str]:
    if current_platform.is_cuda():
        return ["FLASH_ATTN_VLLM_V1", "TRITON_ATTN_VLLM_V1", "TREE_ATTN"]
    elif current_platform.is_rocm():
        attn_backend_list = ["TRITON_ATTN_VLLM_V1"]
        try:
            import aiter  # noqa: F401
            attn_backend_list.append("FLASH_ATTN_VLLM_V1")
        except Exception:
            print("Skip FLASH_ATTN_VLLM_V1 on ROCm as aiter is not installed")

        return attn_backend_list
    else:
        raise ValueError("Unsupported platform")