utils.py 29.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import asyncio
4
import copy
5
import functools
6
import os
7
import signal
8
9
import subprocess
import sys
10
import tempfile
11
import time
12
import warnings
13
from contextlib import contextmanager, suppress
14
from pathlib import Path
15
from typing import Any, Callable, Literal, Optional, Union
16

17
import cloudpickle
18
import openai
19
import pytest
20
import requests
21
import torch
22
import torch.nn.functional as F
23
from openai.types.completion import Completion
24
from typing_extensions import ParamSpec
25

26
import vllm.envs as envs
27
from tests.models.utils import TextTextLogprobs
28
29
from vllm.distributed import (ensure_model_parallel_initialized,
                              init_distributed_environment)
30
from vllm.engine.arg_utils import AsyncEngineArgs
31
from vllm.entrypoints.openai.cli_args import make_arg_parser
32
from vllm.model_executor.model_loader.loader import get_model_loader
33
from vllm.platforms import current_platform
34
from vllm.transformers_utils.tokenizer import get_tokenizer
35
from vllm.utils import (FlexibleArgumentParser, GB_bytes,
36
                        cuda_device_count_stateless, get_open_port)
37

38
if current_platform.is_rocm():
39
40
41
42
43
44
45
46
47
48
49
    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()
50
elif current_platform.is_cuda():
51
52
53
    from vllm.third_party.pynvml import (nvmlDeviceGetHandleByIndex,
                                         nvmlDeviceGetMemoryInfo, nvmlInit,
                                         nvmlShutdown)
54
55
56
57
58
59
60
61

    @contextmanager
    def _nvml():
        try:
            nvmlInit()
            yield
        finally:
            nvmlShutdown()
62
63
64
65
66
else:

    @contextmanager
    def _nvml():
        yield
67

68

69
70
VLLM_PATH = Path(__file__).parent.parent
"""Path to root of the vLLM repository."""
71
72


73
74
class RemoteOpenAIServer:
    DUMMY_API_KEY = "token-abc123"  # vLLM's OpenAI server does not need API key
75
76
77

    def __init__(self,
                 model: str,
78
                 vllm_serve_args: list[str],
79
                 *,
80
                 env_dict: Optional[dict[str, str]] = None,
81
                 seed: Optional[int] = 0,
82
83
                 auto_port: bool = True,
                 max_wait_seconds: Optional[float] = None) -> None:
84
        if auto_port:
85
86
            if "-p" in vllm_serve_args or "--port" in vllm_serve_args:
                raise ValueError("You have manually specified the port "
87
88
                                 "when `auto_port=True`.")

89
90
91
92
            # Don't mutate the input args
            vllm_serve_args = vllm_serve_args + [
                "--port", str(get_open_port())
            ]
93
94
95
96
97
98
        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)]
99

Ethan Xu's avatar
Ethan Xu committed
100
101
102
        parser = FlexibleArgumentParser(
            description="vLLM's remote OpenAI server.")
        parser = make_arg_parser(parser)
103
        args = parser.parse_args(["--model", model, *vllm_serve_args])
104
105
106
        self.host = str(args.host or 'localhost')
        self.port = int(args.port)

107
108
109
110
        # 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)
111
112
113
114
115
            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)
116

117
118
119
120
        env = os.environ.copy()
        # the current process might initialize cuda,
        # to be safe, we should use spawn method
        env['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
121
122
        if env_dict is not None:
            env.update(env_dict)
123
124
125
126
127
128
        self.proc = subprocess.Popen(
            ["vllm", "serve", model, *vllm_serve_args],
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
        )
129
        max_wait_seconds = max_wait_seconds or 240
130
        self._wait_for_server(url=self.url_for("health"),
131
                              timeout=max_wait_seconds)
132
133
134
135
136
137

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.proc.terminate()
138
        try:
139
            self.proc.wait(8)
140
141
142
        except subprocess.TimeoutExpired:
            # force kill if needed
            self.proc.kill()
143
144
145
146
147
148
149
150

    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
151
152
153
154
155
            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`.
156
157
                result = self.proc.poll()
                if result is not None and result != 0:
158
                    raise RuntimeError("Server exited unexpectedly.") from None
159
160
161
162

                time.sleep(0.5)
                if time.time() - start > timeout:
                    raise RuntimeError(
163
                        "Server failed to start in time.") from None
164
165
166
167
168
169
170
171

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

172
173
174
    def get_client(self, **kwargs):
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
175
176
177
        return openai.OpenAI(
            base_url=self.url_for("v1"),
            api_key=self.DUMMY_API_KEY,
178
179
            max_retries=0,
            **kwargs,
180
181
        )

182
    def get_async_client(self, **kwargs):
183
184
        if "timeout" not in kwargs:
            kwargs["timeout"] = 600
185
186
187
188
        return openai.AsyncOpenAI(base_url=self.url_for("v1"),
                                  api_key=self.DUMMY_API_KEY,
                                  max_retries=0,
                                  **kwargs)
189
190


191
192
193
194
def _test_completion(
    client: openai.OpenAI,
    model: str,
    prompt: str,
195
    token_ids: list[int],
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
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
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
):
    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


295
296
297
298
299
300
301
302
303
304
305
306
307
308
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)

309
310
    logprobs = completion.choices[0].logprobs.top_logprobs[0]
    logprobs = {k: round(v, 2) for k, v in logprobs.items()}
311
312
313

    results.append({
        "test": "completion_close",
314
        "logprobs": logprobs,
315
316
317
318
319
    })

    return results


320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
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


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


411
def compare_two_settings(model: str,
412
413
414
415
                         arg1: list[str],
                         arg2: list[str],
                         env1: Optional[dict[str, str]] = None,
                         env2: Optional[dict[str, str]] = None,
416
                         *,
417
                         method: str = "generate",
418
                         max_wait_seconds: Optional[float] = None) -> None:
419
    """
420
421
422
423
424
425
426
427
428
    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.
429
430
    """

431
432
433
434
435
436
437
438
439
440
    compare_all_settings(
        model,
        [arg1, arg2],
        [env1, env2],
        method=method,
        max_wait_seconds=max_wait_seconds,
    )


def compare_all_settings(model: str,
441
442
                         all_args: list[list[str]],
                         all_envs: list[Optional[dict[str, str]]],
443
                         *,
444
                         method: str = "generate",
445
446
447
448
449
450
451
452
453
454
                         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.
    """

455
    trust_remote_code = False
456
    for args in all_args:
457
458
459
460
461
        if "--trust-remote-code" in args:
            trust_remote_code = True
            break

    tokenizer_mode = "auto"
462
    for args in all_args:
463
464
465
466
467
468
469
470
471
        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,
    )
472

473
474
475
476
477
478
479
    can_force_load_format = True

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

480
    prompt = "Hello, my name is"
481
    token_ids = tokenizer(prompt).input_ids
482
    ref_results: list = []
483
    for i, (args, env) in enumerate(zip(all_args, all_envs)):
484
485
486
487
488
489
490
491
492
        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]
493
        compare_results: list = []
494
        results = ref_results if i == 0 else compare_results
495
496
497
498
        with RemoteOpenAIServer(model,
                                args,
                                env_dict=env,
                                max_wait_seconds=max_wait_seconds) as server:
499
500
501
502
503
504
505
506
507
508
509
510
            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,
            })

511
512
            if method == "generate":
                results += _test_completion(client, model, prompt, token_ids)
513
514
515
516
517
518
519
            elif method == "generate_close":
                results += _test_completion_close(client, model, prompt)
            elif method == "generate_with_image":
                results += _test_image_text(
                    client, model,
                    "https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png"
                )
520
521
522
            elif method == "encode":
                results += _test_embeddings(client, model, prompt)
            else:
523
                raise ValueError(f"Unknown method: {method}")
524

525
526
527
528
529
530
531
532
            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):
533
534
535
                    ref_result = copy.deepcopy(ref_result)
                    compare_result = copy.deepcopy(compare_result)
                    if "embedding" in ref_result and method == "encode":
536
537
538
539
540
541
                        sim = F.cosine_similarity(
                            torch.tensor(ref_result["embedding"]),
                            torch.tensor(compare_result["embedding"]),
                            dim=0,
                        )
                        assert sim >= 0.999, (
542
                            f"Embedding for {model=} are not the same.\n"
543
                            f"cosine_similarity={sim}\n")
544
545
                        del ref_result["embedding"]
                        del compare_result["embedding"]
546
547
548
549
550
551
                    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")
552
553


554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
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)


570
def multi_process_parallel(
571
    monkeypatch: pytest.MonkeyPatch,
572
573
    tp_size: int,
    pp_size: int,
574
    test_target: Any,
575
) -> None:
576
577
    import ray

578
579
    # Using ray helps debugging the error when it failed
    # as compared to multiprocessing.
580
581
    # NOTE: We need to set working_dir for distributed tests,
    # otherwise we may get import errors on ray workers
582
583
584
585
586
587
    ray.init(runtime_env={"working_dir": VLLM_PATH})

    distributed_init_port = get_open_port()
    refs = []
    for rank in range(tp_size * pp_size):
        refs.append(
588
589
590
591
592
593
594
            test_target.remote(
                monkeypatch,
                tp_size,
                pp_size,
                rank,
                distributed_init_port,
            ), )
595
596
597
    ray.get(refs)

    ray.shutdown()
598
599
600


@contextmanager
601
def error_on_warning(category: type[Warning] = Warning):
602
603
    """
    Within the scope of this context manager, tests will fail if any warning
604
    of the given category is emitted.
605
606
    """
    with warnings.catch_warnings():
607
        warnings.filterwarnings("error", category=category)
608
609

        yield
610
611


612
613
614
615
616
617
618
619
620
621
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]


622
@_nvml()
623
def wait_for_gpu_memory_to_clear(devices: list[int],
624
625
626
627
                                 threshold_bytes: int,
                                 timeout_s: float = 120) -> None:
    # Use nvml instead of pytorch to reduce measurement error from torch cuda
    # context.
628
    devices = get_physical_device_indices(devices)
629
630
    start_time = time.time()
    while True:
631
632
        output: dict[int, str] = {}
        output_raw: dict[int, float] = {}
633
        for device in devices:
634
            if current_platform.is_rocm():
635
636
637
638
639
640
641
                dev_handle = amdsmi_get_processor_handles()[device]
                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
                gb_used = mem_info["vram_used"] / 2**10
            else:
                dev_handle = nvmlDeviceGetHandleByIndex(device)
                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
                gb_used = mem_info.used / 2**30
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
            output_raw[device] = gb_used
            output[device] = f'{gb_used:.02f}'

        print('gpu memory used (GB): ', end='')
        for k, v in output.items():
            print(f'{k}={v}; ', end='')
        print('')

        dur_s = time.time() - start_time
        if all(v <= (threshold_bytes / 2**30) for v in output_raw.values()):
            print(f'Done waiting for free GPU memory on devices {devices=} '
                  f'({threshold_bytes/2**30=}) {dur_s=:.02f}')
            break

        if dur_s >= timeout_s:
            raise ValueError(f'Memory of devices {devices=} not free after '
                             f'{dur_s=:.02f} ({threshold_bytes/2**30=})')

        time.sleep(5)
661
662


663
664
665
666
667
_P = ParamSpec("_P")


def fork_new_process_for_each_test(
        f: Callable[_P, None]) -> Callable[_P, None]:
668
669
670
    """Decorator to fork a new process for each test function.
    See https://github.com/vllm-project/vllm/issues/7053 for more details.
    """
671
672

    @functools.wraps(f)
673
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
674
675
676
677
678
        # 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()
679
        print(f"Fork a new process to run a test {pid}")
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
        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
697
            old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
698
699
700
            # kill all child processes
            os.killpg(pgid, signal.SIGTERM)
            # restore the signal handler
701
            signal.signal(signal.SIGTERM, old_signal_handler)
702
703
704
705
            assert _exitcode == 0, (f"function {f} failed when called with"
                                    f" args {args} and kwargs {kwargs}")

    return wrapper
706
707


708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
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". 
               If not specified,
               it defaults to "spawn" on ROCm platforms and "fork" otherwise.

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

    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


780
def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
781
782
783
    """
    Get a pytest mark, which skips the test if the GPU doesn't meet
    a minimum memory requirement in GB.
784

785
786
    This can be leveraged via `@large_gpu_test` to skip tests in environments
    without enough resources, or called when filtering tests to run directly.
787
788
    """
    try:
789
        if current_platform.is_cpu() or current_platform.is_openvino():
790
791
792
793
794
795
796
797
798
799
            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

800
    return pytest.mark.skipif(
801
        memory_gb < min_gb,
802
        reason=f"Need at least {min_gb}GB GPU memory to run the test.",
803
804
    )

805
806
807
808
809
810
811
812

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

815
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
816
        return mark(f)
817
818
819
820

    return wrapper


821
822
823
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)
824
825
826
827
828
    test_skipif = pytest.mark.skipif(
        cuda_device_count_stateless() < num_gpus,
        reason=f"Need at least {num_gpus} GPUs to run the test.",
    )

829
830
831
832
833
834
835
836
837
    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)

838
    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
839
        func = create_new_process_for_each_test()(f)
840
841
842
843
        for mark in reversed(marks):
            func = mark(func)

        return func
844
845
846
847

    return wrapper


848
async def completions_with_server_args(
849
    prompts: list[str],
850
    model_name: str,
851
    server_cli_args: list[str],
852
853
    num_logprobs: Optional[int],
    max_wait_seconds: int = 240,
854
    max_tokens: Union[int, list] = 5,
855
) -> list[Completion]:
856
857
858
859
860
861
862
863
864
865
    '''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
866
867
868
      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.
869
870
871
872
873

    Returns:
      OpenAI Completion instance
    '''

874
875
876
877
878
    if isinstance(max_tokens, int):
        max_tokens = [max_tokens] * len(prompts)

    assert len(max_tokens) == len(prompts)

879
880
881
882
883
    outputs = None
    with RemoteOpenAIServer(model_name,
                            server_cli_args,
                            max_wait_seconds=max_wait_seconds) as server:
        client = server.get_async_client()
884
885
886
887
888
889
890
891
892
        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)

893
    assert outputs is not None, "Completion API call failed."
894
895
896
897

    return outputs


898
def get_client_text_generations(completions: list[Completion]) -> list[str]:
899
900
901
    '''Extract generated tokens from the output of a
    request made to an Open-AI-protocol completions endpoint.
    '''
902
903
    assert all([len(x.choices) == 1 for x in completions])
    return [x.choices[0].text for x in completions]
904
905
906


def get_client_text_logprob_generations(
907
        completions: list[Completion]) -> list[TextTextLogprobs]:
908
909
910
911
912
913
914
915
    '''Operates on the output of a request made to an Open-AI-protocol
    completions endpoint; obtains top-rank logprobs for each token in
    each :class:`SequenceGroup`
    '''
    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))
916
            for completion in completions for x in completion.choices]