@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 134 lines 3.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Test pause/resume with Data Parallel (DP) via HTTP API.
4
5 This example demonstrates coordinated pause/resume across multiple DP ranks.
6 The pause synchronizes across all DP engines via all-reduce.
7
8 Prerequisites:
9 Start a vLLM server with data parallelism:
10
11 $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
12 --enforce-eager \
13 --data-parallel-size 4 \
14 --tensor-parallel-size 1
15
16 Then run this script:
17
18 $ python data_parallel_pause_resume.py
19
20 The test verifies pause works by:
21 1. Starting a streaming generation request
22 2. Pausing the server mid-generation
23 3. Sleeping for PAUSE_DURATION seconds
24 4. Resuming the server
25 5. Verifying there was a gap in token generation matching the pause duration
26 """
27
28 import argparse
29 import threading
30 import time
31
32 import requests
33 from openai import OpenAI
34
35 BASE_URL = "http://localhost:8000"
36 MODEL_NAME = "facebook/opt-125m"
37 PAUSE_DURATION = 3.0
38
39
40 def pause_generation(base_url: str, mode: str = "keep") -> None:
41 """Pause generation via HTTP endpoint."""
42 url = f"{base_url}/pause"
43 response = requests.post(url, params={"mode": mode}, timeout=60)
44 response.raise_for_status()
45 print("Server paused")
46
47
48 def resume_generation(base_url: str) -> None:
49 """Resume generation via HTTP endpoint."""
50 url = f"{base_url}/resume"
51 response = requests.post(url, timeout=60)
52 response.raise_for_status()
53 print("Server resumed")
54
55
56 def main():
57 parser = argparse.ArgumentParser()
58 parser.add_argument("--base-url", default=BASE_URL)
59 parser.add_argument("--model", default=MODEL_NAME)
60 args = parser.parse_args()
61
62 client = OpenAI(
63 base_url=f"{args.base_url}/v1",
64 api_key="EMPTY",
65 )
66
67 prompt = "Write a long story about a dragon. Once upon a time"
68 token_times: list[float] = []
69 pause_token_idx = 0
70 pause_triggered = threading.Event()
71
72 def generator_thread():
73 """Stream tokens and record timestamps."""
74 stream = client.completions.create(
75 model=args.model,
76 prompt=prompt,
77 max_tokens=50,
78 stream=True,
79 )
80 for chunk in stream:
81 if chunk.choices[0].text:
82 token_times.append(time.monotonic())
83 token_count = len(token_times)
84 print(f"Token {token_count}: {chunk.choices[0].text!r}")
85
86 # Signal controller after some tokens
87 if token_count >= 5 and not pause_triggered.is_set():
88 pause_triggered.set()
89
90 def controller_thread():
91 """Pause and resume the server."""
92 nonlocal pause_token_idx
93
94 # Wait for some tokens
95 pause_triggered.wait()
96
97 print(f"\nPausing server (keep mode) at token {len(token_times)}...")
98 pause_generation(args.base_url, mode="keep")
99 pause_token_idx = len(token_times)
100 print(f"Sleeping for {PAUSE_DURATION}s...")
101
102 time.sleep(PAUSE_DURATION)
103
104 print("Resuming server...")
105 resume_generation(args.base_url)
106 print("Resumed!\n")
107
108 # Run both threads
109 gen_thread = threading.Thread(target=generator_thread)
110 ctrl_thread = threading.Thread(target=controller_thread)
111
112 gen_thread.start()
113 ctrl_thread.start()
114
115 gen_thread.join()
116 ctrl_thread.join()
117
118 # Check gap at the pause point
119 if pause_token_idx < len(token_times):
120 pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
121 print(
122 f"\nGap after pause (token {pause_token_idx} -> "
123 f"{pause_token_idx + 1}): {pause_gap:.3f}s"
124 )
125 if pause_gap >= PAUSE_DURATION * 0.9:
126 print("Test passed! Pause synchronized across DP ranks.")
127 else:
128 print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
129 else:
130 print("Test failed! No tokens were generated after resuming.")
131
132
133 if __name__ == "__main__":
134 main()