main
py 109 lines 3.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Simple example demonstrating streaming offline inference with AsyncLLM (V1 engine).
4
5 This script shows the core functionality of vLLM's AsyncLLM engine for streaming
6 token-by-token output in offline inference scenarios. It demonstrates DELTA mode
7 streaming where you receive new tokens as they are generated.
8
9 Usage:
10 python examples/deployment/async_llm_streaming.py
11 """
12
13 import asyncio
14
15 from vllm import SamplingParams
16 from vllm.engine.arg_utils import AsyncEngineArgs
17 from vllm.sampling_params import RequestOutputKind
18 from vllm.v1.engine.async_llm import AsyncLLM
19
20
21 async def stream_response(engine: AsyncLLM, prompt: str, request_id: str) -> None:
22 """Stream response from AsyncLLM and display tokens as they arrive.
23
24 This function demonstrates the core streaming pattern:
25 1. Create SamplingParams with DELTA output kind
26 2. Call engine.generate() and iterate over the async generator
27 3. Print new tokens as they arrive
28 4. Handle the finished flag to know when generation is complete
29 """
30 print(f"\n🚀 Prompt: {prompt!r}")
31 print("💬 Response: ", end="", flush=True)
32
33 # Configure sampling parameters for streaming
34 sampling_params = SamplingParams(
35 max_tokens=100,
36 temperature=0.8,
37 top_p=0.95,
38 seed=42, # For reproducible results
39 output_kind=RequestOutputKind.DELTA, # Get only new tokens each iteration
40 )
41
42 try:
43 # Stream tokens from AsyncLLM
44 async for output in engine.generate(
45 request_id=request_id, prompt=prompt, sampling_params=sampling_params
46 ):
47 # Process each completion in the output
48 for completion in output.outputs:
49 # In DELTA mode, we get only new tokens generated since last iteration
50 new_text = completion.text
51 if new_text:
52 print(new_text, end="", flush=True)
53
54 # Check if generation is finished
55 if output.finished:
56 print("\n✅ Generation complete!")
57 break
58
59 except Exception as e:
60 print(f"\n❌ Error during streaming: {e}")
61 raise
62
63
64 async def main():
65 print("🔧 Initializing AsyncLLM...")
66
67 # Create AsyncLLM engine with simple configuration
68 engine_args = AsyncEngineArgs(
69 model="meta-llama/Llama-3.2-1B-Instruct",
70 enforce_eager=True, # Faster startup for examples
71 )
72 engine = AsyncLLM.from_engine_args(engine_args)
73
74 try:
75 # Example prompts to demonstrate streaming
76 prompts = [
77 "The future of artificial intelligence is",
78 "In a galaxy far, far away",
79 "The key to happiness is",
80 ]
81
82 print(f"🎯 Running {len(prompts)} streaming examples...")
83
84 # Process each prompt
85 for i, prompt in enumerate(prompts, 1):
86 print(f"\n{'=' * 60}")
87 print(f"Example {i}/{len(prompts)}")
88 print(f"{'=' * 60}")
89
90 request_id = f"stream-example-{i}"
91 await stream_response(engine, prompt, request_id)
92
93 # Brief pause between examples
94 if i < len(prompts):
95 await asyncio.sleep(0.5)
96
97 print("\n🎉 All streaming examples completed!")
98
99 finally:
100 # Always clean up the engine
101 print("🔧 Shutting down engine...")
102 engine.shutdown()
103
104
105 if __name__ == "__main__":
106 try:
107 asyncio.run(main())
108 except KeyboardInterrupt:
109 print("\n🛑 Interrupted by user")