@@ -144,7 +144,7 @@ def test_engine_core(monkeypatch):
144
144
def test_engine_core_advanced_sampling (monkeypatch ):
145
145
"""
146
146
A basic end-to-end test to verify that the engine functions correctly
147
- when additional sampling parameters, such as min_tokens and
147
+ when additional sampling parameters, such as top_p, min_tokens, and
148
148
presence_penalty, are set.
149
149
"""
150
150
with monkeypatch .context () as m :
@@ -167,11 +167,23 @@ def test_engine_core_advanced_sampling(monkeypatch):
167
167
stop_token_ids = [1001 , 1002 ],
168
168
)
169
169
engine_core .add_request (request )
170
- assert len (engine_core .scheduler .waiting ) == 1
171
- assert len (engine_core .scheduler .running ) == 0
172
- # Loop through until they are all done.
173
- while len (engine_core .step ().outputs ) > 0 :
174
- pass
175
170
176
- assert len (engine_core .scheduler .waiting ) == 0
177
- assert len (engine_core .scheduler .running ) == 0
171
+ def _check_engine_state ():
172
+ assert len (engine_core .scheduler .waiting ) == 1
173
+ assert len (engine_core .scheduler .running ) == 0
174
+ # Loop through until they are all done.
175
+ while len (engine_core .step ().outputs ) > 0 :
176
+ pass
177
+ assert len (engine_core .scheduler .waiting ) == 0
178
+ assert len (engine_core .scheduler .running ) == 0
179
+
180
+ _check_engine_state ()
181
+
182
+ # Second request.
183
+ request2 = make_request ()
184
+ request2 .sampling_params = SamplingParams (
185
+ top_p = 0.99 ,
186
+ top_k = 50 ,
187
+ )
188
+ engine_core .add_request (request2 )
189
+ _check_engine_state ()
0 commit comments