A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 05:00 UTC
[Bugfix] Fix spawn_new_process_for_each_test silently swallowing test failures (#41423)
Signed-off-by: dqzhengAP <dqzheng1996@gmail.com>
David Zheng committed
May 6, 2026 at 08:17 UTC
ee38750a7565ee7158ddc78ff5abd12e9c0d3733
4 files changed
+306
-245
tests/utils.py
+49
-36
@@ -19,7 +19,7 @@ import threading
19
import time
20
import warnings
21
from collections.abc import Callable, Iterable, Sequence
22
-from contextlib import ExitStack, contextmanager, suppress
22
+from contextlib import ExitStack, contextmanager
23
from multiprocessing import Process
24
from pathlib import Path
25
from typing import Any, Literal
@@ -1512,52 +1512,65 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non
1512
1513
1514
def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
1515
- """Decorator to spawn a new process for each test function."""
1515
+ """Decorator to spawn a new process for each test function.
1516
+
1517
+ Uses subprocess with cloudpickle to serialize the test function and
1518
+ propagates exceptions back to the parent, so test failures are never
1519
+ silently swallowed (fixes https://github.com/vllm-project/vllm/issues/41415).
1520
+ """
1521
1522
@functools.wraps(f)
1523
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
1519
- # Check if we're already in a subprocess
1520
- if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
1521
- # If we are, just run the function directly
1522
- return f(*args, **kwargs)
1523
-
1524
- import torch.multiprocessing as mp
1525
-
1526
- with suppress(RuntimeError):
1527
- mp.set_start_method("spawn")
1528
-
1529
- # Get the module
1530
- module_name = f.__module__
1531
-
1532
- # Create a process with environment variable set
1533
- env = os.environ.copy()
1534
- env["RUNNING_IN_SUBPROCESS"] = "1"
1535
-
1536
- with tempfile.TemporaryDirectory() as tempdir:
1537
- output_filepath = os.path.join(tempdir, "new_process.tmp")
1524
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".tb", mode="wb") as tmp:
1525
+ tb_file = tmp.name
1526
1539
- # `cloudpickle` allows pickling complex functions directly
1540
- input_bytes = cloudpickle.dumps((f, output_filepath))
1527
+ try:
1528
+ # Serialize the function + args with cloudpickle so closures work
1529
+ payload = cloudpickle.dumps((f, args, kwargs, tb_file))
1530
+
1531
+ child_script = (
1532
+ "import sys, cloudpickle, traceback\n"
1533
+ "try:\n"
1534
+ " from _pytest.outcomes import Skipped\n"
1535
+ "except ImportError:\n"
1536
+ " class Skipped(BaseException): pass\n"
1537
+ "f, args, kwargs, tb_file = "
1538
+ "cloudpickle.loads(sys.stdin.buffer.read())\n"
1539
+ "try:\n"
1540
+ " f(*args, **kwargs)\n"
1541
+ "except Skipped:\n"
1542
+ " sys.exit(0)\n"
1543
+ "except BaseException:\n"
1544
+ " open(tb_file, 'w').write(traceback.format_exc())\n"
1545
+ " sys.exit(1)\n"
1546
+ )
1547
1548
repo_root = str(VLLM_PATH.resolve())
1543
-
1544
- env = dict(env or os.environ)
1549
+ env = os.environ.copy()
1550
env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
1551
1547
- cmd = [sys.executable, "-m", f"{module_name}"]
1548
-
1549
- returned = subprocess.run(
1550
- cmd, input=input_bytes, capture_output=True, env=env
1552
+ result = subprocess.run(
1553
+ [sys.executable, "-c", child_script],
1554
+ input=payload,
1555
+ capture_output=True,
1556
+ env=env,
1557
)
1558
1553
- # check if the subprocess is successful
1554
- try:
1555
- returned.check_returncode()
1556
- except Exception as e:
1557
- # wrap raised exception to provide more information
1559
+ if result.returncode != 0:
1560
+ # Read traceback written by child, fall back to stderr
1561
+ tb = ""
1562
+ if os.path.exists(tb_file) and os.path.getsize(tb_file) > 0:
1563
+ with open(tb_file) as fp:
1564
+ tb = fp.read()
1565
+ else:
1566
+ tb = result.stderr.decode()
1567
raise RuntimeError(
1559
- f"Error raised in subprocess:\n{returned.stderr.decode()}"
1560
- ) from e
1568
+ f"Test subprocess '{f.__name__}' failed "
1569
+ f"(exit code {result.returncode}):\n{tb}"
1570
+ )
1571
+ finally:
1572
+ with contextlib.suppress(OSError):
1573
+ os.remove(tb_file)
1574
1575
return wrapper
1576
tests/utils_/test_spawn_decorator.py
new
+33
@@ -0,0 +1,33 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for spawn_new_process_for_each_test decorator."""
4
+
5
+import pytest
6
+
7
+from tests.utils import spawn_new_process_for_each_test
8
+
9
+
10
+@spawn_new_process_for_each_test
11
+def test_spawn_decorator_passing():
12
+ """Passing function should complete normally."""
13
+ assert 1 + 1 == 2
14
+
15
+
16
+@pytest.mark.xfail(raises=RuntimeError, strict=True)
17
+@spawn_new_process_for_each_test
18
+def test_spawn_decorator_failure_is_caught():
19
+ """Failing function should raise RuntimeError, never silently pass."""
20
+ raise ValueError("intentional failure")
21
+
22
+
23
+@spawn_new_process_for_each_test
24
+def test_spawn_decorator_skip():
25
+ """pytest.skip inside subprocess should propagate correctly."""
26
+ pytest.skip("intentional skip")
27
+
28
+
29
+@spawn_new_process_for_each_test
30
+@pytest.mark.parametrize("x,y,expected", [(1, 2, 3), (0, 0, 0)])
31
+def test_spawn_decorator_parametrized(x, y, expected):
32
+ """Args and kwargs must be forwarded correctly to subprocess."""
33
+ assert x + y == expected
tests/v1/cudagraph/test_cudagraph_dispatch.py
+186
-179
@@ -371,193 +371,200 @@ class TestCUDAGraphWrapper:
371
assert not wrapper.concrete_cudagraph_entries
372
373
374
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
375
-class TestCudagraphIntegration:
376
- def setup_method(self):
377
- # only FULL mode for non-uniform batches
378
- self.comp_config = CompilationConfig(
379
- mode=CompilationMode.VLLM_COMPILE,
380
- cudagraph_mode="FULL",
381
- cudagraph_capture_sizes=[10, 20],
382
- )
383
- self.vllm_config = _create_vllm_config(self.comp_config)
384
- self.dispatcher = CudagraphDispatcher(self.vllm_config)
385
- self.dispatcher.initialize_cudagraph_keys(
386
- self.comp_config.cudagraph_mode, uniform_decode_query_len=1
387
- )
388
-
389
- def _run_and_monitor_call(
390
- self, wrapper, input_tensor, runtime_mode, batch_descriptor
374
+def _run_and_monitor_call(
375
+ wrapper, input_tensor, runtime_mode, batch_descriptor, vllm_config
376
+):
377
+ """Helper to run a single call and monitor the action."""
378
+
379
+ with (
380
+ patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context,
381
+ patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable,
382
):
392
- """Helper to run a single call and monitor the action."""
393
-
394
- with (
395
- patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context,
396
- patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable,
397
- ):
398
- entry = wrapper.concrete_cudagraph_entries.get(batch_descriptor, None)
399
-
400
- context = set_forward_context(
401
- attn_metadata=None,
402
- vllm_config=self.vllm_config,
403
- cudagraph_runtime_mode=runtime_mode,
404
- batch_descriptor=batch_descriptor,
405
- )
406
- mock_replay = MagicMock()
407
- if entry and entry.cudagraph:
408
- with (
409
- context,
410
- patch.object(
411
- entry.cudagraph, "replay", new_callable=MagicMock
412
- ) as mock_replay,
413
- ):
414
- wrapper(input_tensor)
415
- else:
416
- with context:
417
- wrapper(input_tensor)
418
-
419
- if mock_graph_context.called:
420
- # note that this is globally mocked, so it will be detected
421
- # even whether called by the inner or outer wrapper
422
- return "capture_global"
423
- if mock_replay.called:
424
- # only for outer wrapper
425
- return "replay"
426
- if mock_runnable.call_count > 0:
427
- # only for outer wrapper
428
- return "bypass"
429
- return "unknown"
430
-
431
- @create_new_process_for_each_test("spawn")
432
- def test_capture_replay_bypass_logic(self):
433
- model = SimpleMLP().to(DEVICE_TYPE)
434
- full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
435
- max_bs = 16
436
- persistent_input_buffer = torch.zeros(max_bs, 10, device=DEVICE_TYPE)
437
- input_1 = persistent_input_buffer[:1]
438
- input_2 = persistent_input_buffer[:2]
439
- input_3 = persistent_input_buffer[:3]
440
-
441
- desc_1 = BatchDescriptor(num_tokens=1)
442
- desc_2 = BatchDescriptor(num_tokens=2)
443
- desc_3_unseen = BatchDescriptor(num_tokens=3)
383
+ entry = wrapper.concrete_cudagraph_entries.get(batch_descriptor, None)
384
445
- # 0. global warmup
446
- with set_forward_context(
385
+ context = set_forward_context(
386
attn_metadata=None,
448
- vllm_config=self.vllm_config,
449
- cudagraph_runtime_mode=CUDAGraphMode.NONE,
450
- batch_descriptor=None,
451
- ):
452
- full_wrapper(input_1)
387
+ vllm_config=vllm_config,
388
+ cudagraph_runtime_mode=runtime_mode,
389
+ batch_descriptor=batch_descriptor,
390
+ )
391
+ mock_replay = MagicMock()
392
+ if entry and entry.cudagraph:
393
+ with (
394
+ context,
395
+ patch.object(
396
+ entry.cudagraph, "replay", new_callable=MagicMock
397
+ ) as mock_replay,
398
+ ):
399
+ wrapper(input_tensor)
400
+ else:
401
+ with context:
402
+ wrapper(input_tensor)
403
+
404
+ if mock_graph_context.called:
405
+ # note that this is globally mocked, so it will be detected
406
+ # even whether called by the inner or outer wrapper
407
+ return "capture_global"
408
+ if mock_replay.called:
409
+ # only for outer wrapper
410
+ return "replay"
411
+ if mock_runnable.call_count > 0:
412
+ # only for outer wrapper
413
+ return "bypass"
414
+ return "unknown"
415
+
416
+
417
+@create_new_process_for_each_test("spawn")
418
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
419
+def test_capture_replay_bypass_logic():
420
+ comp_config = CompilationConfig(
421
+ mode=CompilationMode.VLLM_COMPILE,
422
+ cudagraph_mode="FULL",
423
+ cudagraph_capture_sizes=[1, 2],
424
+ )
425
+ vllm_config = _create_vllm_config(comp_config)
426
+ dispatcher = CudagraphDispatcher(vllm_config)
427
+ dispatcher.initialize_cudagraph_keys(
428
+ comp_config.cudagraph_mode, uniform_decode_query_len=1
429
+ )
430
+ model = SimpleMLP().to(DEVICE_TYPE)
431
+ full_wrapper = CUDAGraphWrapper(model, vllm_config, CUDAGraphMode.FULL)
432
+ max_bs = 16
433
+ persistent_input_buffer = torch.zeros(max_bs, 10, device=DEVICE_TYPE)
434
+ input_1 = persistent_input_buffer[:1]
435
+ input_2 = persistent_input_buffer[:2]
436
+ input_3 = persistent_input_buffer[:3]
437
+
438
+ desc_1 = BatchDescriptor(num_tokens=1)
439
+ desc_2 = BatchDescriptor(num_tokens=2)
440
+ desc_3_unseen = BatchDescriptor(num_tokens=3)
441
+
442
+ # 0. global warmup
443
+ with set_forward_context(
444
+ attn_metadata=None,
445
+ vllm_config=vllm_config,
446
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
447
+ batch_descriptor=None,
448
+ ):
449
+ full_wrapper(input_1)
450
454
- rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_1.num_tokens)
455
- # 1. Capture first shape
456
- action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)
457
- assert action == "capture_global"
451
+ rt_mode, key = dispatcher.dispatch(num_tokens=desc_1.num_tokens)
452
+ # 1. Capture first shape
453
+ action = _run_and_monitor_call(full_wrapper, input_1, rt_mode, key, vllm_config)
454
+ assert action == "capture_global"
455
459
- # 2. Replay first shape
460
- action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)
461
- assert action == "replay"
456
+ # 2. Replay first shape
457
+ action = _run_and_monitor_call(full_wrapper, input_1, rt_mode, key, vllm_config)
458
+ assert action == "replay"
459
463
- rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_2.num_tokens)
464
- # 3. Capture second shape
465
- action = self._run_and_monitor_call(full_wrapper, input_2, rt_mode, key)
466
- assert action == "capture_global"
460
+ rt_mode, key = dispatcher.dispatch(num_tokens=desc_2.num_tokens)
461
+ # 3. Capture second shape
462
+ action = _run_and_monitor_call(full_wrapper, input_2, rt_mode, key, vllm_config)
463
+ assert action == "capture_global"
464
468
- # 4. Replay second shape
469
- action = self._run_and_monitor_call(
470
- full_wrapper, input_2, CUDAGraphMode.FULL, desc_2
465
+ # 4. Replay second shape
466
+ action = _run_and_monitor_call(
467
+ full_wrapper, input_2, CUDAGraphMode.FULL, key, vllm_config
468
+ )
469
+ assert action == "replay"
470
+
471
+ # 5. Bypass if no key match
472
+ rt_mode, key = dispatcher.dispatch(num_tokens=desc_3_unseen.num_tokens)
473
+ assert rt_mode == CUDAGraphMode.NONE
474
+ action = _run_and_monitor_call(full_wrapper, input_3, rt_mode, key, vllm_config)
475
+ assert action == "bypass"
476
+
477
+ # capture unseen shape is not allowed after disable
478
+ set_cudagraph_capturing_enabled(False)
479
+ with pytest.raises(RuntimeError):
480
+ _run_and_monitor_call(
481
+ full_wrapper, input_3, CUDAGraphMode.FULL, desc_3_unseen, vllm_config
482
)
472
- assert action == "replay"
483
+ set_cudagraph_capturing_enabled(True)
484
474
- # 5. Bypass if no key match
475
- rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_3_unseen.num_tokens)
476
- assert rt_mode == CUDAGraphMode.NONE
477
- action = self._run_and_monitor_call(full_wrapper, input_3, rt_mode, key)
478
- assert action == "bypass"
479
-
480
- # capture unseen shape is not allowed after disable
481
- set_cudagraph_capturing_enabled(False)
482
- with pytest.raises(RuntimeError):
483
- self._run_and_monitor_call(
484
- full_wrapper, input_3, CUDAGraphMode.FULL, desc_3_unseen
485
- )
486
- set_cudagraph_capturing_enabled(True)
487
-
488
- @create_new_process_for_each_test("spawn")
489
- def test_nested_wrappers(self):
490
- """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
491
- model = SimpleMLP().to(DEVICE_TYPE)
492
- full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
493
- input_1 = torch.randn(1, 10, device=DEVICE_TYPE)
494
-
495
- # Setup: Inner model is wrapped with PIECEWISE, outer with FULL
496
- inner_model = SimpleMLP().to(DEVICE_TYPE)
497
- piecewise_wrapper = CUDAGraphWrapper(
498
- inner_model, self.vllm_config, CUDAGraphMode.PIECEWISE
499
- )
500
- inner_model.forward = MagicMock(wraps=inner_model.forward)
501
- outer_model = SimpleMLP().to(DEVICE_TYPE)
502
- # When outer model is called, it calls the piecewise_wrapper
503
- outer_model.forward = MagicMock(
504
- wraps=outer_model.forward, side_effect=piecewise_wrapper
505
- )
506
- full_wrapper = CUDAGraphWrapper(
507
- outer_model, self.vllm_config, CUDAGraphMode.FULL
508
- )
485
510
- desc_1 = BatchDescriptor(num_tokens=1)
486
+@create_new_process_for_each_test("spawn")
487
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
488
+def test_nested_wrappers():
489
+ """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
490
+ comp_config = CompilationConfig(
491
+ mode=CompilationMode.VLLM_COMPILE,
492
+ cudagraph_mode="FULL",
493
+ cudagraph_capture_sizes=[1],
494
+ )
495
+ vllm_config = _create_vllm_config(comp_config)
496
+ dispatcher = CudagraphDispatcher(vllm_config)
497
+ dispatcher.initialize_cudagraph_keys(
498
+ comp_config.cudagraph_mode, uniform_decode_query_len=1
499
+ )
500
+ model = SimpleMLP().to(DEVICE_TYPE)
501
+ full_wrapper = CUDAGraphWrapper(model, vllm_config, CUDAGraphMode.FULL)
502
+ input_1 = torch.randn(1, 10, device=DEVICE_TYPE)
503
+
504
+ # Setup: Inner model is wrapped with PIECEWISE, outer with FULL
505
+ inner_model = SimpleMLP().to(DEVICE_TYPE)
506
+ piecewise_wrapper = CUDAGraphWrapper(
507
+ inner_model, vllm_config, CUDAGraphMode.PIECEWISE
508
+ )
509
+ inner_model.forward = MagicMock(wraps=inner_model.forward)
510
+ outer_model = SimpleMLP().to(DEVICE_TYPE)
511
+ # When outer model is called, it calls the piecewise_wrapper
512
+ outer_model.forward = MagicMock(
513
+ wraps=outer_model.forward, side_effect=piecewise_wrapper
514
+ )
515
+ full_wrapper = CUDAGraphWrapper(outer_model, vllm_config, CUDAGraphMode.FULL)
516
512
- # 0. global warmup
513
- with set_forward_context(
514
- attn_metadata=None,
515
- vllm_config=self.vllm_config,
516
- cudagraph_runtime_mode=CUDAGraphMode.NONE,
517
- batch_descriptor=None,
518
- ):
519
- full_wrapper(input_1)
520
-
521
- # --- Test runtime mode FULL---
522
- # Run with FULL mode context. Expect outer wrapper to capture.
523
- # The inner mock should be called once inside the graph capture.
524
- outer_model.forward.reset_mock()
525
- inner_model.forward.reset_mock()
526
- action = self._run_and_monitor_call(
527
- full_wrapper, input_1, CUDAGraphMode.FULL, desc_1
528
- )
529
- assert action == "capture_global"
530
- assert outer_model.forward.call_count == 1
531
- assert inner_model.forward.call_count == 1
532
-
533
- # Run again. Expect outer wrapper to replay.
534
- # The outer model should NOT be called because the whole graph
535
- # is replayed.
536
- action = self._run_and_monitor_call(
537
- full_wrapper, input_1, CUDAGraphMode.FULL, desc_1
538
- )
539
- assert action == "replay"
540
- assert outer_model.forward.call_count == 1 # No new call
541
- assert inner_model.forward.call_count == 1
542
-
543
- # --- Test runtime mode PIECEWISE ---
544
- outer_model.forward.reset_mock()
545
- inner_model.forward.reset_mock()
546
- # Run with PIECEWISE mode context.
547
- # Expect outer wrapper to bypass and call inner wrapper.
548
- # Inner wrapper should capture.
549
- action = self._run_and_monitor_call(
550
- full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1
551
- )
552
- assert action == "capture_global"
553
- assert outer_model.forward.call_count == 1
554
- assert inner_model.forward.call_count == 1
555
-
556
- # Run again with PIECEWISE.
557
- # Outer bypasses, inner replays.
558
- action = self._run_and_monitor_call(
559
- full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1
560
- )
561
- assert action == "bypass"
562
- assert outer_model.forward.call_count == 2
563
- assert inner_model.forward.call_count == 1
517
+ desc_1 = BatchDescriptor(num_tokens=1)
518
+
519
+ # 0. global warmup
520
+ with set_forward_context(
521
+ attn_metadata=None,
522
+ vllm_config=vllm_config,
523
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
524
+ batch_descriptor=None,
525
+ ):
526
+ full_wrapper(input_1)
527
+
528
+ # --- Test runtime mode FULL---
529
+ # Run with FULL mode context. Expect outer wrapper to capture.
530
+ # The inner mock should be called once inside the graph capture.
531
+ outer_model.forward.reset_mock()
532
+ inner_model.forward.reset_mock()
533
+ action = _run_and_monitor_call(
534
+ full_wrapper, input_1, CUDAGraphMode.FULL, desc_1, vllm_config
535
+ )
536
+ assert action == "capture_global"
537
+ assert outer_model.forward.call_count == 1
538
+ assert inner_model.forward.call_count == 1
539
+
540
+ # Run again. Expect outer wrapper to replay.
541
+ # The outer model should NOT be called because the whole graph
542
+ # is replayed.
543
+ action = _run_and_monitor_call(
544
+ full_wrapper, input_1, CUDAGraphMode.FULL, desc_1, vllm_config
545
+ )
546
+ assert action == "replay"
547
+ assert outer_model.forward.call_count == 1 # No new call
548
+ assert inner_model.forward.call_count == 1
549
+
550
+ # --- Test runtime mode PIECEWISE ---
551
+ outer_model.forward.reset_mock()
552
+ inner_model.forward.reset_mock()
553
+ # Run with PIECEWISE mode context.
554
+ # Expect outer wrapper to bypass and call inner wrapper.
555
+ # Inner wrapper should capture.
556
+ action = _run_and_monitor_call(
557
+ full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1, vllm_config
558
+ )
559
+ assert action == "capture_global"
560
+ assert outer_model.forward.call_count == 1
561
+ assert inner_model.forward.call_count == 1
562
+
563
+ # Run again with PIECEWISE.
564
+ # Outer bypasses, inner replays.
565
+ action = _run_and_monitor_call(
566
+ full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1, vllm_config
567
+ )
568
+ assert action == "bypass"
569
+ assert outer_model.forward.call_count == 2
570
+ assert inner_model.forward.call_count == 1
tests/v1/logits_processors/test_custom_online.py
+38
-30
@@ -120,12 +120,11 @@ api_keyword_args = {
120
121
122
@create_new_process_for_each_test()
123
-@pytest.mark.asyncio
123
@pytest.mark.parametrize(
124
"model_name",
125
[MODEL_NAME],
126
)
128
-async def test_custom_logitsprocs(client: openai.AsyncOpenAI, model_name: str):
127
+def test_custom_logitsprocs(server, model_name: str):
128
"""Test custom logitsprocs when starting OpenAI server from CLI
129
130
Launch vLLM OpenAI-compatible server, configured to load a custom logitproc
@@ -139,36 +138,45 @@ async def test_custom_logitsprocs(client: openai.AsyncOpenAI, model_name: str):
138
token
139
"""
140
142
- use_dummy_logitproc = True
143
- for prompt in prompts:
144
- # Build request arguments
145
- request_keyword_args: dict[str, Any] = {
146
- **api_keyword_args,
147
- }
148
- if use_dummy_logitproc:
149
- # 50% of requests pass target_token custom arg
150
- target_token = random.choice([128, 67])
151
- # For requests which activate the dummy logitproc, choose one of
152
- # two `target_token` values which are known not to be EOS tokens
153
- request_keyword_args["extra_body"] = {
154
- "vllm_xargs": {DUMMY_LOGITPROC_ARG: target_token}
155
- }
156
- batch = await client.completions.create(
157
- model=model_name,
158
- prompt=prompt,
159
- **request_keyword_args,
160
- )
141
+ import asyncio
142
162
- if use_dummy_logitproc:
163
- # Only for requests which activate dummy logitproc - validate that
164
- # output token is repeated
165
- choices: openai.types.CompletionChoice = batch.choices
166
- toks = choices[0].logprobs.tokens
167
- if not all([x == toks[0] for x in toks]):
168
- raise AssertionError(f"Generated {toks} should all be {toks[0]}")
143
+ async def _async_main(srv, mn):
144
+ async with srv.get_async_client() as client:
145
+ await _run(client)
146
170
- # Alternate whether to activate dummy logitproc for each request
171
- use_dummy_logitproc = not use_dummy_logitproc
147
+ async def _run(client):
148
+ use_dummy_logitproc = True
149
+ for prompt in prompts:
150
+ # Build request arguments
151
+ request_keyword_args: dict[str, Any] = {
152
+ **api_keyword_args,
153
+ }
154
+ if use_dummy_logitproc:
155
+ # 50% of requests pass target_token custom arg
156
+ target_token = random.choice([128, 67])
157
+ # For requests which activate the dummy logitproc, choose one of
158
+ # two `target_token` values which are known not to be EOS tokens
159
+ request_keyword_args["extra_body"] = {
160
+ "vllm_xargs": {DUMMY_LOGITPROC_ARG: target_token}
161
+ }
162
+ batch = await client.completions.create(
163
+ model=model_name,
164
+ prompt=prompt,
165
+ **request_keyword_args,
166
+ )
167
+
168
+ if use_dummy_logitproc:
169
+ # Only for requests which activate dummy logitproc - validate that
170
+ # output token is repeated
171
+ choices: openai.types.CompletionChoice = batch.choices
172
+ toks = choices[0].logprobs.tokens
173
+ if not all([x == toks[0] for x in toks]):
174
+ raise AssertionError(f"Generated {toks} should all be {toks[0]}")
175
+
176
+ # Alternate whether to activate dummy logitproc for each request
177
+ use_dummy_logitproc = not use_dummy_logitproc
178
+
179
+ asyncio.run(_async_main(server, model_name))
180
181
182
@pytest.mark.asyncio