@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 172 lines 5.52 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import asyncio
5 from collections.abc import AsyncGenerator
6 from unittest.mock import AsyncMock, MagicMock
7
8 import pytest
9
10 from vllm.engine.protocol import StreamingInput
11 from vllm.outputs import RequestOutput
12 from vllm.sampling_params import RequestOutputKind, SamplingParams
13 from vllm.v1.engine.async_llm import AsyncLLM
14 from vllm.v1.engine.output_processor import RequestOutputCollector
15
16
17 @pytest.fixture
18 def mock_async_llm():
19 """Create a mock AsyncLLM with mocked dependencies."""
20 # Create a minimal mock without initializing the full engine
21 llm = MagicMock(spec=AsyncLLM)
22
23 # Mock the essential attributes
24 llm.vllm_config = MagicMock()
25 llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
26 llm.model_config = MagicMock()
27 llm.model_config.max_model_len = 2048
28 llm.log_requests = False
29 llm.errored = False
30 llm._pause_cond = asyncio.Condition()
31 llm._paused = False
32
33 # Mock methods
34 llm._run_output_handler = MagicMock()
35 llm.abort = AsyncMock()
36
37 # Use the real generate method from AsyncLLM
38 llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
39
40 return llm
41
42
43 @pytest.mark.asyncio
44 async def test_generate_normal_flow(mock_async_llm):
45 """Test normal generation flow with streaming requests."""
46 request_id = "test_request"
47 prompt = "Tell me about Paris"
48 sampling_params = SamplingParams(max_tokens=10)
49
50 # Create a mock queue with outputs
51 queue = RequestOutputCollector(RequestOutputKind.FINAL_ONLY, request_id)
52 output1 = RequestOutput(
53 request_id=request_id,
54 prompt="Tell me about Paris",
55 prompt_token_ids=[1, 2, 3],
56 prompt_logprobs=None,
57 outputs=[],
58 finished=False,
59 )
60 output2 = RequestOutput(
61 request_id=request_id,
62 prompt="Tell me about Paris",
63 prompt_token_ids=[1, 2, 3],
64 prompt_logprobs=None,
65 outputs=[],
66 finished=True,
67 )
68
69 # Feed outputs to queue as they're consumed to avoid aggregation
70 async def feed_outputs():
71 queue.put(output1)
72 await asyncio.sleep(1) # Let first output be consumed
73 queue.put(output2)
74
75 asyncio.create_task(feed_outputs()) # noqa
76
77 # Mock add_request to return the queue
78 async def mock_add_request(*args, **kwargs):
79 return queue
80
81 mock_async_llm.add_request = mock_add_request
82
83 # Collect outputs from generate
84 outputs = []
85 async for output in mock_async_llm.generate(
86 prompt=prompt,
87 sampling_params=sampling_params,
88 request_id=request_id,
89 ):
90 outputs.append(output)
91
92 assert len(outputs) == 2
93 assert outputs[0].finished is False
94 assert outputs[1].finished is True
95
96
97 def make_output(request_id: str, finished: bool) -> RequestOutput:
98 """Helper to create a RequestOutput."""
99 return RequestOutput(
100 request_id=request_id,
101 prompt="test",
102 prompt_token_ids=[1, 2, 3],
103 prompt_logprobs=None,
104 outputs=[],
105 finished=finished,
106 )
107
108
109 @pytest.mark.asyncio
110 async def test_generate_with_async_generator():
111 """Test generate with an async input generator.
112
113 With the new streaming input API, completion is signaled by finishing
114 the input generator (not via a resumable flag). Each input chunk
115 produces intermediate outputs, and the final output has finished=True.
116 """
117 request_id = "test"
118 sampling_params = SamplingParams(max_tokens=10)
119
120 llm = MagicMock(spec=AsyncLLM)
121 llm.vllm_config = MagicMock()
122 llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
123 llm.model_config = MagicMock()
124 llm.model_config.max_model_len = 2048
125 llm.log_requests = False
126 llm.errored = False
127 llm._pause_cond = asyncio.Condition()
128 llm._paused = False
129 llm._run_output_handler = MagicMock()
130 llm.abort = AsyncMock()
131
132 # Bind the real generate method
133 llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
134
135 # Track inputs processed
136 inputs_received = []
137 queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id)
138
139 async def mock_add_request(req_id, prompt, params, *args, **kwargs):
140 # When prompt is an AsyncGenerator, process streaming inputs
141 if isinstance(prompt, AsyncGenerator):
142 # Process inputs in background, produce outputs
143 async def handle_stream():
144 async for input_chunk in prompt:
145 inputs_received.append(input_chunk.prompt)
146 # Each input produces an intermediate output
147 queue.put(make_output(req_id, finished=False))
148 await asyncio.sleep(0.01)
149 # Final output when stream ends
150 queue.put(make_output(req_id, finished=True))
151
152 asyncio.create_task(handle_stream())
153 return queue
154 return queue
155
156 llm.add_request = mock_add_request
157
158 async def input_generator() -> AsyncGenerator[StreamingInput, None]:
159 yield StreamingInput(prompt="Hello", sampling_params=sampling_params)
160 yield StreamingInput(prompt=" world", sampling_params=sampling_params)
161
162 outputs = []
163 async for output in llm.generate(input_generator(), sampling_params, request_id):
164 outputs.append(output)
165
166 # Two intermediate outputs + one final output
167 assert len(outputs) == 3
168 assert outputs[0].finished is False
169 assert outputs[1].finished is False
170 assert outputs[2].finished is True
171 # Both inputs were processed
172 assert inputs_received == ["Hello", " world"]