@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e19bce40a1

Fork 0

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 14, 05:30 UTC

[V0 Deprecation] Remove AsyncLLMEngine (#25025)

Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Sep 18, 2025 at 11:07 UTC e19bce40a1660cb7c03b790d0b000db155cf925d

11 files changed +78 -1969

tests/entrypoints/openai/test_chat.py

+19 -35

@@ -28,11 +28,9 @@ def monkeypatch_module():

28 mpatch.undo()

31 -@pytest.fixture(scope="module", params=[False, True])

32 -def server(request, monkeypatch_module, zephyr_lora_files): #noqa: F811

33 -

34 - use_v1 = request.param

35 - monkeypatch_module.setenv('VLLM_USE_V1', '1' if use_v1 else '0')

31 +@pytest.fixture(scope="module")

32 +def server(monkeypatch_module, zephyr_lora_files): #noqa: F811

33 + monkeypatch_module.setenv('VLLM_USE_V1', '1')

35 args = [

36 # use half precision for speed and memory savings in CI environment

@@ -57,13 +55,6 @@ def server(request, monkeypatch_module, zephyr_lora_files): #noqa: F811

55 yield remote_server

60 -@pytest.fixture

61 -def is_v1_server(server):

62 - import os

63 - assert os.environ['VLLM_USE_V1'] in ['0', '1']

64 - return os.environ['VLLM_USE_V1'] == '1'

65 -

66 -

58 @pytest_asyncio.fixture

59 async def client(server):

60 async with server.get_async_client() as async_client:

@@ -481,10 +472,9 @@ async def test_chat_completion_stream_options(client: openai.AsyncOpenAI,

472

473 @pytest.mark.asyncio

474 async def test_structured_outputs_choice_chat(

484 - client: openai.AsyncOpenAI, sample_structured_outputs_choices,

485 - is_v1_server: bool):

486 - if not is_v1_server:

487 - pytest.skip("Structured outputs is only supported in v1 engine")

475 + client: openai.AsyncOpenAI,

476 + sample_structured_outputs_choices,

477 +):

478 messages = [{

479 "role": "system",

480 "content": "you are a helpful assistant"

@@ -522,12 +512,10 @@ async def test_structured_outputs_choice_chat(

512

513

514 @pytest.mark.asyncio

525 -async def test_structured_outputs_json_chat(client: openai.AsyncOpenAI,

526 - sample_json_schema,

527 - is_v1_server: bool):

528 - if not is_v1_server:

529 - pytest.skip("Structured outputs is only supported in v1 engine")

530 -

515 +async def test_structured_outputs_json_chat(

516 + client: openai.AsyncOpenAI,

517 + sample_json_schema,

518 +):

519 messages = [{

520 "role": "system",

521 "content": "you are a helpful assistant"

@@ -569,10 +557,10 @@ async def test_structured_outputs_json_chat(client: openai.AsyncOpenAI,

557

558

559 @pytest.mark.asyncio

572 -async def test_structured_outputs_regex_chat(client: openai.AsyncOpenAI,

573 - sample_regex, is_v1_server: bool):

574 - if not is_v1_server:

575 - pytest.skip("Structured outputs is only supported in v1 engine")

560 +async def test_structured_outputs_regex_chat(

561 + client: openai.AsyncOpenAI,

562 + sample_regex,

563 +):

564

565 messages = [{

566 "role": "system",

@@ -660,10 +648,10 @@ async def test_structured_outputs_choice_chat_logprobs(

648

649

650 @pytest.mark.asyncio

663 -async def test_named_tool_use(client: openai.AsyncOpenAI, sample_json_schema,

664 - is_v1_server: bool):

665 - if not is_v1_server:

666 - pytest.skip("Tool use is only supported in v1 engine")

651 +async def test_named_tool_use(

652 + client: openai.AsyncOpenAI,

653 + sample_json_schema,

654 +):

655 messages = [{

656 "role": "system",

657 "content": "you are a helpful assistant"

@@ -821,11 +809,7 @@ async def test_response_format_json_object(client: openai.AsyncOpenAI):

809

810

811 @pytest.mark.asyncio

824 -async def test_response_format_json_schema(client: openai.AsyncOpenAI,

825 - is_v1_server: bool):

826 - if not is_v1_server:

827 - pytest.skip(

828 - "JSON schema response format is only supported in v1 engine")

812 +async def test_response_format_json_schema(client: openai.AsyncOpenAI):

813 prompt = 'what is 1+1? The format is "result": 2'

814 # Check that this prompt cannot lead to a valid JSON without json_schema

815 for _ in range(2):

tests/entrypoints/openai/test_completion.py deleted

-830

@@ -1,830 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -# imports for structured outputs tests

4 -import json

5 -import os

6 -from typing import Optional

7 -

8 -import jsonschema

9 -import openai # use the official client for correctness check

10 -import pytest

11 -import pytest_asyncio

12 -import regex as re

13 -import requests

14 -# downloading lora to test lora requests

15 -from openai import BadRequestError

16 -

17 -from vllm.transformers_utils.tokenizer import get_tokenizer

18 -

19 -from ...utils import RemoteOpenAIServer

20 -

21 -# any model with a chat template should work here

22 -MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

23 -# technically these adapters use a different base model,

24 -# but we're not testing generation quality here

25 -

26 -

27 -@pytest.fixture(scope="module")

28 -def default_server_args(zephyr_lora_files):

29 - return [

30 - # use half precision for speed and memory savings in CI environment

31 - "--dtype",

32 - "bfloat16",

33 - "--max-model-len",

34 - "8192",

35 - "--max-num-seqs",

36 - "128",

37 - "--enforce-eager",

38 - # lora config

39 - "--enable-lora",

40 - "--lora-modules",

41 - f"zephyr-lora={zephyr_lora_files}",

42 - "--max-lora-rank",

43 - "64",

44 - "--max-cpu-loras",

45 - "2",

46 - ]

47 -

48 -

49 -@pytest.fixture(scope="module",

50 - params=["", "--disable-frontend-multiprocessing"])

51 -def server(default_server_args, request):

52 - if request.param:

53 - default_server_args.append(request.param)

54 -

55 - original_value = os.environ.get('VLLM_USE_V1')

56 - os.environ['VLLM_USE_V1'] = '0'

57 - try:

58 - with RemoteOpenAIServer(MODEL_NAME,

59 - default_server_args) as remote_server:

60 - yield remote_server

61 - finally:

62 - # Restore original env value

63 - if original_value is None:

64 - os.environ.pop('VLLM_USE_V1', None)

65 - else:

66 - os.environ['VLLM_USE_V1'] = original_value

67 -

68 -

69 -@pytest.fixture

70 -def is_v1_server(server):

71 - import os

72 -

73 - # For completion tests, we assume v0 since there's no explicit v1 setup

74 - return os.environ.get('VLLM_USE_V1', '0') == '1'

75 -

76 -

77 -@pytest_asyncio.fixture

78 -async def client(server):

79 - async with server.get_async_client() as async_client:

80 - yield async_client

81 -

82 -

83 -@pytest.mark.asyncio

84 -@pytest.mark.parametrize(

85 - # first test base model, then test loras

86 - "model_name",

87 - [MODEL_NAME, "zephyr-lora"],

88 -)

89 -async def test_single_completion(client: openai.AsyncOpenAI, model_name: str):

90 - completion = await client.completions.create(model=model_name,

91 - prompt="Hello, my name is",

92 - max_tokens=5,

93 - temperature=0.0)

94 -

95 - assert completion.id is not None

96 - assert completion.choices is not None and len(completion.choices) == 1

97 -

98 - choice = completion.choices[0]

99 - assert len(choice.text) >= 5

100 - assert choice.finish_reason == "length"

101 - assert completion.usage == openai.types.CompletionUsage(

102 - completion_tokens=5, prompt_tokens=6, total_tokens=11)

103 -

104 - # test using token IDs

105 - completion = await client.completions.create(

106 - model=model_name,

107 - prompt=[0, 0, 0, 0, 0],

108 - max_tokens=5,

109 - temperature=0.0,

110 - )

111 - assert len(completion.choices[0].text) >= 1

112 - assert completion.choices[0].prompt_logprobs is None

113 -

114 -

115 -@pytest.mark.asyncio

116 -async def test_added_lora_tokens_base_model(client: openai.AsyncOpenAI):

117 - # test using token IDs

118 - with pytest.raises(openai.BadRequestError, match="out of vocabulary"):

119 - # Added tokens should be rejected by the base model

120 - await client.completions.create(

121 - model=MODEL_NAME,

122 - prompt=[0, 0, 32000, 32001, 32002],

123 - echo=True,

124 - max_tokens=5,

125 - temperature=0.0,

126 - )

127 -

128 -

129 -@pytest.mark.asyncio

130 -@pytest.mark.parametrize(

131 - # first test base model, then test loras

132 - "model_name",

133 - [MODEL_NAME, "zephyr-lora"],

134 -)

135 -async def test_no_logprobs(client: openai.AsyncOpenAI, model_name: str):

136 - # test using token IDs

137 - completion = await client.completions.create(

138 - model=model_name,

139 - prompt=[0, 0, 0, 0, 0],

140 - max_tokens=5,

141 - temperature=0.0,

142 - logprobs=None,

143 - )

144 - choice = completion.choices[0]

145 - assert choice.logprobs is None

146 -

147 -

148 -@pytest.mark.asyncio

149 -@pytest.mark.parametrize(

150 - # just test 1 lora

151 - "model_name",

152 - [MODEL_NAME, "zephyr-lora"],

153 -)

154 -async def test_zero_logprobs(client: openai.AsyncOpenAI, model_name: str):

155 - # test using token IDs

156 - completion = await client.completions.create(

157 - model=model_name,

158 - prompt=[0, 0, 0, 0, 0],

159 - max_tokens=5,

160 - temperature=0.0,

161 - logprobs=0,

162 - )

163 - choice = completion.choices[0]

164 - assert choice.logprobs is not None

165 - assert choice.logprobs.token_logprobs is not None

166 - assert choice.logprobs.top_logprobs is not None

167 - assert len(choice.logprobs.top_logprobs[0]) == 1

168 -

169 -

170 -@pytest.mark.asyncio

171 -@pytest.mark.parametrize(

172 - "model_name",

173 - [MODEL_NAME, "zephyr-lora"],

174 -)

175 -async def test_some_logprobs(client: openai.AsyncOpenAI, model_name: str):

176 - # test using token IDs

177 - completion = await client.completions.create(

178 - model=model_name,

179 - prompt=[0, 0, 0, 0, 0],

180 - max_tokens=5,

181 - temperature=0.0,

182 - logprobs=5,

183 - )

184 - choice = completion.choices[0]

185 - assert choice.logprobs is not None

186 - assert choice.logprobs.token_logprobs is not None

187 - assert choice.logprobs.top_logprobs is not None

188 - assert 5 <= len(choice.logprobs.top_logprobs[0]) <= 6

189 -

190 -

191 -@pytest.mark.asyncio

192 -@pytest.mark.parametrize(

193 - "model_name",

194 - [MODEL_NAME, "zephyr-lora"],

195 -)

196 -async def test_too_many_completion_logprobs(client: openai.AsyncOpenAI,

197 - model_name: str):

198 -

199 - with pytest.raises(

200 - (openai.BadRequestError, openai.APIError)): # test using token IDs

201 - await client.completions.create(

202 - model=model_name,

203 - prompt=[0, 0, 0, 0, 0],

204 - max_tokens=5,

205 - temperature=0.0,

206 - # vLLM has higher default max_logprobs (20 instead of 5) to support

207 - # both Completion API and Chat Completion API

208 - logprobs=21,

209 - )

210 - ...

211 - with pytest.raises(

212 - (openai.BadRequestError, openai.APIError)): # test using token IDs

213 - stream = await client.completions.create(

214 - model=model_name,

215 - prompt=[0, 0, 0, 0, 0],

216 - max_tokens=5,

217 - temperature=0.0,

218 - # vLLM has higher default max_logprobs (20 instead of 5) to support

219 - # both Completion API and Chat Completion API

220 - logprobs=30,

221 - stream=True,

222 - )

223 - async for chunk in stream:

224 - ...

225 -

226 - # the server should still work afterwards

227 - completion = await client.completions.create(

228 - model=model_name,

229 - prompt=[0, 0, 0, 0, 0],

230 - max_tokens=5,

231 - temperature=0.0,

232 - )

233 - assert len(completion.choices[0].text) >= 0

234 -

235 -

236 -@pytest.mark.asyncio

237 -@pytest.mark.parametrize("model_name, prompt_logprobs", [(MODEL_NAME, -1),

238 - (MODEL_NAME, 0),

239 - (MODEL_NAME, 1),

240 - (MODEL_NAME, None)])

241 -async def test_prompt_logprobs_completion(client: openai.AsyncOpenAI,

242 - model_name: str,

243 - prompt_logprobs: Optional[int]):

244 - params: dict = {

245 - "prompt": ["A robot may not injure another robot", "My name is"],

246 - "model": model_name,

247 - }

248 - if prompt_logprobs is not None:

249 - params["extra_body"] = {"prompt_logprobs": prompt_logprobs}

250 -

251 - if prompt_logprobs is not None and prompt_logprobs < 0:

252 - with pytest.raises(BadRequestError):

253 - await client.completions.create(**params)

254 - else:

255 - completion = await client.completions.create(**params)

256 - if prompt_logprobs is not None:

257 - assert completion.choices[0].prompt_logprobs is not None

258 - assert len(completion.choices[0].prompt_logprobs) > 0

259 -

260 - assert completion.choices[1].prompt_logprobs is not None

261 - assert len(completion.choices[1].prompt_logprobs) > 0

262 -

263 - else:

264 - assert completion.choices[0].prompt_logprobs is None

265 -

266 -

267 -@pytest.mark.asyncio

268 -@pytest.mark.parametrize(

269 - "model_name",

270 - [MODEL_NAME, "zephyr-lora"],

271 -)

272 -async def test_completion_streaming(client: openai.AsyncOpenAI,

273 - model_name: str):

274 - prompt = "What is an LLM?"

275 -

276 - single_completion = await client.completions.create(

277 - model=model_name,

278 - prompt=prompt,

279 - max_tokens=5,

280 - temperature=0.0,

281 - )

282 - single_output = single_completion.choices[0].text

283 - stream = await client.completions.create(model=model_name,

284 - prompt=prompt,

285 - max_tokens=5,

286 - temperature=0.0,

287 - stream=True)

288 - chunks: list[str] = []

289 - finish_reason_count = 0

290 - async for chunk in stream:

291 - chunks.append(chunk.choices[0].text)

292 - if chunk.choices[0].finish_reason is not None:

293 - finish_reason_count += 1

294 - # finish reason should only return in last block

295 - assert finish_reason_count == 1

296 - assert chunk.choices[0].finish_reason == "length"

297 - assert chunk.choices[0].text

298 - assert "".join(chunks) == single_output

299 -

300 -

301 -@pytest.mark.asyncio

302 -@pytest.mark.parametrize(

303 - "model_name",

304 - [MODEL_NAME, "zephyr-lora"],

305 -)

306 -async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):

307 - """Streaming for parallel sampling.

308 - The tokens from multiple samples, are flattened into a single stream,

309 - with an index to indicate which sample the token belongs to.

310 - """

311 -

312 - prompt = "What is an LLM?"

313 - n = 3

314 - max_tokens = 5

315 -

316 - stream = await client.completions.create(model=model_name,

317 - prompt=prompt,

318 - max_tokens=max_tokens,

319 - n=n,

320 - stream=True)

321 - chunks: list[list[str]] = [[] for i in range(n)]

322 - finish_reason_count = 0

323 - async for chunk in stream:

324 - index = chunk.choices[0].index

325 - text = chunk.choices[0].text

326 - chunks[index].append(text)

327 - if chunk.choices[0].finish_reason is not None:

328 - finish_reason_count += 1

329 - assert finish_reason_count == n

330 - for chunk in chunks:

331 - assert len(chunk) == max_tokens

332 - print("".join(chunk))

333 -

334 -

335 -@pytest.mark.asyncio

336 -@pytest.mark.parametrize(

337 - "model_name",

338 - [MODEL_NAME, "zephyr-lora"],

339 -)

340 -async def test_completion_stream_options(client: openai.AsyncOpenAI,

341 - model_name: str):

342 - prompt = "What is the capital of France?"

343 -

344 - # Test stream=True, stream_options=

345 - # {"include_usage": False, "continuous_usage_stats": False}

346 - stream = await client.completions.create(model=model_name,

347 - prompt=prompt,

348 - max_tokens=5,

349 - temperature=0.0,

350 - stream=True,

351 - stream_options={

352 - "include_usage": False,

353 - "continuous_usage_stats":

354 - False,

355 - })

356 -

357 - async for chunk in stream:

358 - assert chunk.usage is None

359 -

360 - # Test stream=True, stream_options=

361 - # {"include_usage": False, "continuous_usage_stats": True}

362 - stream = await client.completions.create(model=model_name,

363 - prompt=prompt,

364 - max_tokens=5,

365 - temperature=0.0,

366 - stream=True,

367 - stream_options={

368 - "include_usage": False,

369 - "continuous_usage_stats":

370 - True,

371 - })

372 - async for chunk in stream:

373 - assert chunk.usage is None

374 -

375 - # Test stream=True, stream_options=

376 - # {"include_usage": True, "continuous_usage_stats": False}

377 - stream = await client.completions.create(model=model_name,

378 - prompt=prompt,

379 - max_tokens=5,

380 - temperature=0.0,

381 - stream=True,

382 - stream_options={

383 - "include_usage": True,

384 - "continuous_usage_stats":

385 - False,

386 - })

387 - async for chunk in stream:

388 - if chunk.choices[0].finish_reason is None:

389 - assert chunk.usage is None

390 - else:

391 - assert chunk.usage is None

392 - final_chunk = await stream.__anext__()

393 - assert final_chunk.usage is not None

394 - assert final_chunk.usage.prompt_tokens > 0

395 - assert final_chunk.usage.completion_tokens > 0

396 - assert final_chunk.usage.total_tokens == (

397 - final_chunk.usage.prompt_tokens +

398 - final_chunk.usage.completion_tokens)

399 - assert final_chunk.choices == []

400 -

401 - # Test stream=True, stream_options=

402 - # {"include_usage": True, "continuous_usage_stats": True}

403 - stream = await client.completions.create(model=model_name,

404 - prompt=prompt,

405 - max_tokens=5,

406 - temperature=0.0,

407 - stream=True,

408 - stream_options={

409 - "include_usage": True,

410 - "continuous_usage_stats":

411 - True,

412 - })

413 - async for chunk in stream:

414 - assert chunk.usage is not None

415 - assert chunk.usage.prompt_tokens > 0

416 - assert chunk.usage.completion_tokens > 0

417 - assert chunk.usage.total_tokens == (chunk.usage.prompt_tokens +

418 - chunk.usage.completion_tokens)

419 - if chunk.choices[0].finish_reason is not None:

420 - final_chunk = await stream.__anext__()

421 - assert final_chunk.usage is not None

422 - assert final_chunk.usage.prompt_tokens > 0

423 - assert final_chunk.usage.completion_tokens > 0

424 - assert final_chunk.usage.total_tokens == (

425 - final_chunk.usage.prompt_tokens +

426 - final_chunk.usage.completion_tokens)

427 - assert final_chunk.choices == []

428 -

429 - # Test stream=False, stream_options=

430 - # {"include_usage": None}

431 - with pytest.raises(BadRequestError):

432 - await client.completions.create(model=model_name,

433 - prompt=prompt,

434 - max_tokens=5,

435 - temperature=0.0,

436 - stream=False,

437 - stream_options={"include_usage": None})

438 -

439 - # Test stream=False, stream_options=

440 - # {"include_usage": True}

441 - with pytest.raises(BadRequestError):

442 - await client.completions.create(model=model_name,

443 - prompt=prompt,

444 - max_tokens=5,

445 - temperature=0.0,

446 - stream=False,

447 - stream_options={"include_usage": True})

448 -

449 - # Test stream=False, stream_options=

450 - # {"continuous_usage_stats": None}

451 - with pytest.raises(BadRequestError):

452 - await client.completions.create(

453 - model=model_name,

454 - prompt=prompt,

455 - max_tokens=5,

456 - temperature=0.0,

457 - stream=False,

458 - stream_options={"continuous_usage_stats": None})

459 -

460 - # Test stream=False, stream_options=

461 - # {"continuous_usage_stats": True}

462 - with pytest.raises(BadRequestError):

463 - await client.completions.create(

464 - model=model_name,

465 - prompt=prompt,

466 - max_tokens=5,

467 - temperature=0.0,

468 - stream=False,

469 - stream_options={"continuous_usage_stats": True})

470 -

471 -

472 -@pytest.mark.asyncio

473 -@pytest.mark.parametrize(

474 - "model_name",

475 - [MODEL_NAME, "zephyr-lora"],

476 -)

477 -async def test_batch_completions(client: openai.AsyncOpenAI, model_name: str):

478 - # test both text and token IDs

479 - for prompts in (["Hello, my name is"] * 2, [[0, 0, 0, 0, 0]] * 2):

480 - # test simple list

481 - batch = await client.completions.create(

482 - model=model_name,

483 - prompt=prompts,

484 - max_tokens=5,

485 - temperature=0.0,

486 - )

487 - assert len(batch.choices) == 2

488 - assert batch.choices[0].text == batch.choices[1].text

489 -

490 - # test n = 2

491 - batch = await client.completions.create(

492 - model=model_name,

493 - prompt=prompts,

494 - n=2,

495 - max_tokens=5,

496 - temperature=0.0,

497 - extra_body=dict(

498 - # NOTE: this has to be true for n > 1 in vLLM, but

499 - # not necessary for official client.

500 - use_beam_search=True),

501 - )

502 - assert len(batch.choices) == 4

503 - assert batch.choices[0].text != batch.choices[

504 - 1].text, "beam search should be different"

505 - assert batch.choices[0].text == batch.choices[

506 - 2].text, "two copies of the same prompt should be the same"

507 - assert batch.choices[1].text == batch.choices[

508 - 3].text, "two copies of the same prompt should be the same"

509 -

510 - # test streaming

511 - batch = await client.completions.create(

512 - model=model_name,

513 - prompt=prompts,

514 - max_tokens=5,

515 - temperature=0.0,

516 - stream=True,

517 - )

518 - texts = [""] * 2

519 - async for chunk in batch:

520 - assert len(chunk.choices) == 1

521 - choice = chunk.choices[0]

522 - texts[choice.index] += choice.text

523 - assert texts[0] == texts[1]

524 -

525 -

526 -@pytest.mark.asyncio

527 -async def test_logits_bias(client: openai.AsyncOpenAI):

528 - prompt = "Hello, my name is"

529 - max_tokens = 5

530 - tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)

531 -

532 - # Test exclusive selection

533 - token_id = 1000

534 - completion = await client.completions.create(

535 - model=MODEL_NAME,

536 - prompt=prompt,

537 - max_tokens=max_tokens,

538 - temperature=0.0,

539 - logit_bias={str(token_id): 100},

540 - seed=42,

541 - )

542 - assert len(completion.choices[0].text) >= 5

543 - response_tokens = tokenizer(completion.choices[0].text,

544 - add_special_tokens=False)["input_ids"]

545 - expected_tokens = tokenizer(tokenizer.decode([token_id] * 5),

546 - add_special_tokens=False)["input_ids"]

547 - assert all([

548 - response == expected

549 - for response, expected in zip(response_tokens, expected_tokens)

550 - ])

551 -

552 - # Test ban

553 - completion = await client.completions.create(

554 - model=MODEL_NAME,

555 - prompt=prompt,

556 - max_tokens=max_tokens,

557 - temperature=0.0,

558 - )

559 - response_tokens = tokenizer(completion.choices[0].text,

560 - add_special_tokens=False)["input_ids"]

561 - first_response = completion.choices[0].text

562 - completion = await client.completions.create(

563 - model=MODEL_NAME,

564 - prompt=prompt,

565 - max_tokens=max_tokens,

566 - temperature=0.0,

567 - logit_bias={str(token): -100

568 - for token in response_tokens},

569 - )

570 - assert first_response != completion.choices[0].text

571 -

572 -

573 -@pytest.mark.asyncio

574 -async def test_allowed_token_ids(client: openai.AsyncOpenAI):

575 - prompt = "Hello, my name is"

576 - max_tokens = 1

577 - tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)

578 -

579 - # Test exclusive selection

580 - allowed_ids = [21555, 21557, 21558]

581 - completion = await client.completions.create(

582 - model=MODEL_NAME,

583 - prompt=prompt,

584 - max_tokens=max_tokens,

585 - temperature=0.0,

586 - seed=42,

587 - extra_body=dict(allowed_token_ids=allowed_ids),

588 - logprobs=1,

589 - )

590 - response_tokens = completion.choices[0].logprobs.tokens

591 - assert len(response_tokens) == 1

592 - assert tokenizer.convert_tokens_to_ids(response_tokens)[0] in allowed_ids

593 -

594 -

595 -@pytest.mark.asyncio

596 -async def test_structured_outputs_json_completion(

597 - client: openai.AsyncOpenAI,

598 - sample_json_schema,

599 - is_v1_server: bool,

600 -):

601 - if not is_v1_server:

602 - pytest.skip("structured outputs is only supported in v1 engine")

603 -

604 - completion = await client.completions.create(

605 - model=MODEL_NAME,

606 - prompt=f"Give an example JSON for an employee profile "

607 - f"that fits this schema: {sample_json_schema}",

608 - n=3,

609 - temperature=1.0,

610 - max_tokens=500,

611 - extra_body=dict(structured_outputs=dict(json=sample_json_schema)))

612 -

613 - assert completion.id is not None

614 - assert len(completion.choices) == 3

615 - for i in range(3):

616 - output_json = json.loads(completion.choices[i].text)

617 - jsonschema.validate(instance=output_json, schema=sample_json_schema)

618 -

619 -

620 -@pytest.mark.asyncio

621 -async def test_structured_outputs_regex_completion(

622 - client: openai.AsyncOpenAI,

623 - sample_regex,

624 - is_v1_server: bool,

625 -):

626 - if not is_v1_server:

627 - pytest.skip("structured outputs is only supported in v1 engine")

628 -

629 - completion = await client.completions.create(

630 - model=MODEL_NAME,

631 - prompt=f"Give an example IPv4 address with this regex: {sample_regex}",

632 - n=3,

633 - temperature=1.0,

634 - max_tokens=20,

635 - extra_body=dict(structured_outputs=dict(regex=sample_regex)))

636 -

637 - assert completion.id is not None

638 - assert len(completion.choices) == 3

639 - for i in range(3):

640 - assert re.fullmatch(sample_regex,

641 - completion.choices[i].text) is not None

642 -

643 -

644 -@pytest.mark.asyncio

645 -async def test_structured_outputs_choice_completion(

646 - client: openai.AsyncOpenAI,

647 - sample_structured_outputs_choices,

648 - is_v1_server: bool,

649 -):

650 - if not is_v1_server:

651 - pytest.skip("structured outputs is only supported in v1 engine")

652 -

653 - completion = await client.completions.create(

654 - model=MODEL_NAME,

655 - prompt="The best language for type-safe systems programming is ",

656 - n=2,

657 - temperature=1.0,

658 - max_tokens=10,

659 - extra_body=dict(structured_outputs=dict(

660 - choice=sample_structured_outputs_choices)))

661 -

662 - assert completion.id is not None

663 - assert len(completion.choices) == 2

664 - for i in range(2):

665 - assert completion.choices[i].text in sample_structured_outputs_choices

666 -

667 -

668 -@pytest.mark.asyncio

669 -async def test_structured_outputs_grammar(client: openai.AsyncOpenAI,

670 - sample_sql_statements,

671 - is_v1_server: bool):

672 - if not is_v1_server:

673 - pytest.skip("grammar is only supported in v1 engine")

674 -

675 - completion = await client.completions.create(

676 - model=MODEL_NAME,

677 - prompt=("Generate a sql state that select col_1 from "

678 - "table_1 where it is equals to 1"),

679 - temperature=1.0,

680 - max_tokens=500,

681 - extra_body=dict(

682 - structured_outputs=dict(grammar=sample_sql_statements), ))

683 -

684 - content = completion.choices[0].text

685 -

686 - # use Lark to parse the output, and make sure it's a valid parse tree

687 - from lark import Lark

688 - parser = Lark(sample_sql_statements)

689 - parser.parse(content)

690 -

691 - # remove spaces for comparison b/c we removed them in the grammar

692 - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")

693 -

694 - assert content.strip() == ground_truth

695 -

696 -

697 -@pytest.mark.asyncio

698 -@pytest.mark.parametrize(

699 - # first test base model, then test loras

700 - "model_name",

701 - [MODEL_NAME, "zephyr-lora"],

702 -)

703 -@pytest.mark.parametrize("logprobs_arg", [1, 0])

704 -async def test_echo_logprob_completion(client: openai.AsyncOpenAI,

705 - model_name: str, logprobs_arg: int):

706 - tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)

707 - # test using text and token IDs

708 - for prompt in ("Hello, my name is", [0, 0, 0, 0, 0]):

709 - completion = await client.completions.create(model=model_name,

710 - prompt=prompt,

711 - max_tokens=5,

712 - temperature=0.0,

713 - echo=True,

714 - logprobs=logprobs_arg)

715 -

716 - prompt_text = tokenizer.decode(prompt) if isinstance(prompt,

717 - list) else prompt

718 - assert re.search(r"^" + prompt_text, completion.choices[0].text)

719 - logprobs = completion.choices[0].logprobs

720 - assert logprobs is not None

721 - assert len(logprobs.text_offset) > 5

722 - assert (len(logprobs.token_logprobs) > 5

723 - and logprobs.token_logprobs[0] is None)

724 - assert (len(logprobs.top_logprobs) > 5

725 - and logprobs.top_logprobs[0] is None)

726 - for top_logprobs in logprobs.top_logprobs[1:]:

727 - assert max(logprobs_arg,

728 - 1) <= len(top_logprobs) <= logprobs_arg + 1

729 - assert len(logprobs.tokens) > 5

730 -

731 -

732 -@pytest.mark.asyncio

733 -async def test_structured_outputs_type_error(client: openai.AsyncOpenAI,

734 - sample_json_schema, sample_regex,

735 - is_v1_server: bool):

736 - if not is_v1_server:

737 - pytest.skip("structured outputs is only supported in v1 engine")

738 -

739 - with pytest.raises(openai.BadRequestError):

740 - _ = await client.completions.create(

741 - model=MODEL_NAME,

742 - prompt="Give an example JSON that fits this schema: 42",

743 - extra_body=dict(structured_outputs=dict(json=42)))

744 -

745 - with pytest.raises(openai.BadRequestError):

746 - _ = await client.completions.create(

747 - model=MODEL_NAME,

748 - prompt="Give an example string that fits this regex",

749 - extra_body=dict(structured_outputs=dict(

750 - regex=sample_regex,

751 - json=sample_json_schema,

752 - )))

753 -

754 -

755 -@pytest.mark.asyncio

756 -@pytest.mark.parametrize(

757 - "model_name,stream,echo",

758 - [

759 - (MODEL_NAME, False, False),

760 - (MODEL_NAME, False, True),

761 - (MODEL_NAME, True, False),

762 - (MODEL_NAME, True, True) # should not raise BadRequestError error

763 - ],

764 -)

765 -async def test_echo_stream_completion(client: openai.AsyncOpenAI,

766 - model_name: str, stream: bool,

767 - echo: bool):

768 - saying: str = "Hello, my name is"

769 - result = await client.completions.create(model=model_name,

770 - prompt=saying,

771 - max_tokens=10,

772 - temperature=0.0,

773 - echo=echo,

774 - stream=stream)

775 -

776 - stop_reason = "length"

777 -

778 - if not stream:

779 - completion = result

780 - assert completion.id is not None

781 - assert completion.choices is not None and len(completion.choices) == 1

782 -

783 - choice = completion.choices[0]

784 - assert len(choice.text) >= 5

785 - assert choice.finish_reason == stop_reason

786 -

787 - if echo:

788 - assert choice.text is not None and saying in choice.text

789 - else:

790 - assert choice.text is not None and saying not in choice.text

791 -

792 - else:

793 - chunks: list[str] = []

794 - final_finish_reason = None

795 - async for chunk in result:

796 - if chunk.choices and chunk.choices[0].text:

797 - chunks.append(chunk.choices[0].text)

798 - if chunk.choices and chunk.choices[0].finish_reason:

799 - final_finish_reason = chunk.choices[0].finish_reason

800 -

801 - assert final_finish_reason == stop_reason

802 - content = "".join(chunks)

803 - if echo:

804 - assert content is not None and saying in content

805 - else:

806 - assert content is not None and saying not in content

807 -

808 -

809 -@pytest.mark.asyncio

810 -async def test_invocations(server: RemoteOpenAIServer,

811 - client: openai.AsyncOpenAI):

812 - request_args = {

813 - "model": MODEL_NAME,

814 - "prompt": "Hello, my name is",

815 - "max_tokens": 5,

816 - "temperature": 0.0,

817 - "logprobs": None,

818 - }

819 -

820 - completion = await client.completions.create(**request_args)

821 -

822 - invocation_response = requests.post(server.url_for("invocations"),

823 - json=request_args)

824 - invocation_response.raise_for_status()

825 -

826 - completion_output = completion.model_dump()

827 - invocation_output = invocation_response.json()

828 -

829 - assert completion_output.keys() == invocation_output.keys()

830 - assert completion_output["choices"] == invocation_output["choices"]

tests/entrypoints/openai/test_completion_with_prompt_embeds.py

@@ -14,6 +14,9 @@ from transformers import AutoConfig

15 from ...utils import RemoteOpenAIServer

17 +pytest.skip("Skipping prompt_embeds test until V1 supports it.",

18 + allow_module_level=True)

19 +

20 # any model with a chat template should work here

21 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

tests/entrypoints/openai/test_lora_adapters.py

+3 -2

@@ -53,12 +53,13 @@ def monkeypatch_module():

53 mpatch.undo()

56 -@pytest.fixture(scope="module", params=[False, True])

56 +@pytest.fixture(scope="module", params=[True])

57 def server_with_lora_modules_json(request, monkeypatch_module,

58 zephyr_lora_files):

60 use_v1 = request.param

61 - monkeypatch_module.setenv('VLLM_USE_V1', '1' if use_v1 else '0')

61 + assert use_v1

62 + monkeypatch_module.setenv('VLLM_USE_V1', '1')

64 # Define the json format LoRA module configurations

65 lora_module_1 = {

tests/entrypoints/openai/test_metrics.py

+1 -1

@@ -22,7 +22,7 @@ MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

22 PREV_MINOR_VERSION = version._prev_minor_version()

25 -@pytest.fixture(scope="module", params=[True, False])

25 +@pytest.fixture(scope="module", params=[True])

26 def use_v1(request):

27 # Module-scoped variant of run_with_both_engines

28 #

tests/entrypoints/openai/test_return_tokens_as_ids.py

+24 -2

@@ -10,8 +10,30 @@ import pytest

10 from vllm.transformers_utils.tokenizer import get_tokenizer

12 from ...utils import RemoteOpenAIServer

13 -from .test_completion import default_server_args # noqa: F401

14 -from .test_completion import MODEL_NAME

13 +

14 +MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

15 +

16 +

17 +@pytest.fixture(scope="module")

18 +def default_server_args(zephyr_lora_files):

19 + return [

20 + # use half precision for speed and memory savings in CI environment

21 + "--dtype",

22 + "bfloat16",

23 + "--max-model-len",

24 + "8192",

25 + "--max-num-seqs",

26 + "128",

27 + "--enforce-eager",

28 + # lora config

29 + "--enable-lora",

30 + "--lora-modules",

31 + f"zephyr-lora={zephyr_lora_files}",

32 + "--max-lora-rank",

33 + "64",

34 + "--max-cpu-loras",

35 + "2",

36 + ]

39 @pytest.fixture(scope="module")

tests/entrypoints/openai/test_skip_tokenizer.py

-8

@@ -15,14 +15,6 @@ MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"

15 DTYPE = "float16"

18 -@pytest.fixture(autouse=True)

19 -def v1(run_with_both_engines):

20 - # Simple autouse wrapper to run both engines for each test

21 - # This can be promoted up to conftest.py to run for every

22 - # test in a package

23 - pass

24 -

25 -

18 @pytest.fixture(scope="module")

19 def server():

20 args = [

tests/v1/test_oracle.py

-18

@@ -7,7 +7,6 @@ import pytest

7 import vllm.envs as envs

8 from vllm import LLM

9 from vllm.engine.arg_utils import AsyncEngineArgs

10 -from vllm.engine.async_llm_engine import AsyncLLMEngine

11 MODEL = "meta-llama/Llama-3.2-1B-Instruct"

@@ -96,20 +95,3 @@ def test_v1_attn_backend(monkeypatch):

95 _ = AsyncEngineArgs(model=MODEL).create_engine_config()

96 assert envs.VLLM_USE_V1

97 m.delenv("VLLM_USE_V1")

99 -

100 -

101 -def test_reject_using_constructor_directly(monkeypatch):

102 - with monkeypatch.context() as m:

103 - if os.getenv("VLLM_USE_V1", None):

104 - m.delenv("VLLM_USE_V1")

105 -

106 - # Sets VLLM_USE_V1=1.

107 - vllm_config = AsyncEngineArgs(model=MODEL).create_engine_config()

108 -

109 - # This uses the V0 constructor directly.

110 - with pytest.raises(ValueError):

111 - AsyncLLMEngine(vllm_config,

112 - AsyncLLMEngine._get_executor_cls(vllm_config),

113 - log_stats=True)

114 -

115 - m.delenv("VLLM_USE_V1")

vllm/engine/async_llm_engine.py

+2 -1028

@@ -1,1032 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import asyncio

5 -import time

6 -import weakref

7 -from functools import partial

8 -from typing import (Any, AsyncGenerator, Callable, Dict, Iterable, List,

9 - Mapping, Optional, Set, Tuple, Type, Union)

10 -from weakref import ReferenceType

4 +from vllm.v1.engine.async_llm import AsyncLLM

12 -import vllm.envs as envs

13 -from vllm.config import (LoRAConfig, ModelConfig, ParallelConfig,

14 - SchedulerConfig, VllmConfig)

15 -from vllm.core.scheduler import SchedulerOutputs

16 -from vllm.engine.arg_utils import AsyncEngineArgs

17 -from vllm.engine.async_timeout import asyncio_timeout

18 -from vllm.engine.llm_engine import LLMEngine

19 -from vllm.engine.metrics_types import StatLoggerBase

20 -from vllm.engine.protocol import EngineClient

21 -from vllm.executor.executor_base import ExecutorBase

22 -from vllm.inputs import PromptType

23 -from vllm.inputs.preprocess import InputPreprocessor

24 -from vllm.logger import init_logger

25 -from vllm.lora.request import LoRARequest

26 -from vllm.model_executor.layers.sampler import SamplerOutput

27 -from vllm.outputs import PoolingRequestOutput, RequestOutput

28 -from vllm.pooling_params import PoolingParams

29 -from vllm.sampling_params import SamplingParams

30 -from vllm.sequence import ExecuteModelRequest

31 -from vllm.transformers_utils.tokenizer import AnyTokenizer

32 -from vllm.usage.usage_lib import UsageContext

33 -from vllm.utils import Device, deprecate_kwargs, weak_bind

34 -

35 -logger = init_logger(__name__)

36 -ENGINE_ITERATION_TIMEOUT_S = envs.VLLM_ENGINE_ITERATION_TIMEOUT_S

37 -

38 -

39 -class AsyncEngineDeadError(RuntimeError):

40 - pass

41 -

42 -

43 -def _log_task_completion(task: asyncio.Task,

44 - error_callback: Callable[[Exception], None]) -> None:

45 - """This function is only intended for the `engine.run_engine_loop()` task.

46 -

47 - In particular, that task runs a `while True` loop that can only exit if

48 - there is an exception.

49 - """

50 -

51 - exception = None

52 - try:

53 - return_value = task.result()

54 - raise AssertionError(

55 - f"The engine background task should never finish without an "

56 - f"exception. {return_value}")

57 - except asyncio.exceptions.CancelledError:

58 - # We assume that if the task is cancelled, we are gracefully shutting

59 - # down. This should only happen on program exit.

60 - logger.info("Engine is gracefully shutting down.")

61 - except Exception as e:

62 - exception = e

63 - logger.error("Engine background task failed", exc_info=e)

64 - error_callback(exception)

65 - raise AsyncEngineDeadError(

66 - "Task finished unexpectedly. This should never happen! "

67 - "Please open an issue on GitHub. See stack trace above for the "

68 - "actual cause.") from e

69 -

70 -

71 -STOP_ITERATION = Exception() # Sentinel

72 -

73 -

74 -class AsyncStream:

75 - """A stream of RequestOutputs for a request that can be iterated over

76 - asynchronously via an async generator."""

77 -

78 - def __init__(self, request_id: str, cancel: Callable[[str], None]) -> None:

79 - self.request_id = request_id

80 - self._cancel = cancel

81 - self._queue: asyncio.Queue = asyncio.Queue()

82 - self._finished = False

83 -

84 - def put(self, item: Union[RequestOutput, Exception]) -> None:

85 - if not self._finished:

86 - self._queue.put_nowait(item)

87 -

88 - def finish(

89 - self,

90 - exception: Optional[Union[BaseException, Type[BaseException]]] = None,

91 - ) -> None:

92 - if not self._finished:

93 - self._finished = True

94 - self._queue.put_nowait(

95 - exception if self._is_raisable(exception) else STOP_ITERATION)

96 -

97 - @property

98 - def finished(self) -> bool:

99 - return self._finished

100 -

101 - async def generator(self) -> AsyncGenerator[RequestOutput, None]:

102 - try:

103 - while True:

104 - result = await self._queue.get()

105 - if self._is_raisable(result):

106 - if result == STOP_ITERATION:

107 - return

108 - raise result

109 - yield result

110 - except GeneratorExit:

111 - self._cancel(self.request_id)

112 - raise asyncio.CancelledError from None

113 -

114 - @staticmethod

115 - def _is_raisable(value: Any):

116 - return isinstance(value, BaseException) or \

117 - (isinstance(value, type) and \

118 - issubclass(value, BaseException))

119 -

120 -

121 -class RequestTracker:

122 - """Synchronous abstraction for tracking requests."""

123 -

124 - def __init__(self) -> None:

125 - self._request_streams: Dict[str, AsyncStream] = {}

126 - self._aborted_requests: asyncio.Queue[str] = asyncio.Queue()

127 - self._new_requests: asyncio.Queue[Tuple[AsyncStream,

128 - dict]] = asyncio.Queue()

129 - self.new_requests_event = asyncio.Event()

130 -

131 - def __contains__(self, item):

132 - return item in self._request_streams

133 -

134 - def __len__(self) -> int:

135 - return len(self._request_streams)

136 -

137 - def propagate_exception(self,

138 - exc: Exception,

139 - request_id: Optional[str] = None) -> None:

140 - """Propagate an exception to request streams

141 - (all if request_id is None)."""

142 - if request_id is not None:

143 - self.abort_request(request_id, exception=exc)

144 - else:

145 - # NB: tuple() used here because self.abort_request pops the stream

146 - # out of self._request_streams, so we can't iterate on it directly

147 - for rid in tuple(self._request_streams.keys()):

148 - self.abort_request(rid, exception=exc)

149 -

150 - def process_request_output(self,

151 - request_output: RequestOutput,

152 - *,

153 - verbose: bool = False) -> None:

154 - """Process a request output from the engine."""

155 - request_id = request_output.request_id

156 - finished = request_output.finished

157 -

158 - if finished:

159 - stream = self._request_streams.pop(request_id, None)

160 - else:

161 - stream = self._request_streams.get(request_id)

162 - # Guard against a KeyError which can occur if the request was aborted

163 - # while the output was generated

164 - if stream is not None:

165 - stream.put(request_output)

166 - if finished:

167 - stream.finish()

168 -

169 - if verbose and finished:

170 - logger.info("Finished request %s.", request_id)

171 -

172 - def process_exception(self,

173 - request_id: str,

174 - exception: BaseException,

175 - *,

176 - verbose: bool = False) -> None:

177 - """Propagate an exception from the engine."""

178 - if verbose:

179 - logger.info("Finished request %s.", request_id)

180 - self.abort_request(request_id, exception=exception)

181 -

182 - def add_request(self,

183 - request_id: str,

184 - *,

185 - verbose: bool = False,

186 - **engine_add_request_kwargs) -> AsyncStream:

187 - """Add a request to be sent to the engine on the next background

188 - loop iteration."""

189 - if request_id in self._request_streams:

190 - raise KeyError(f"Request {request_id} already exists.")

191 -

192 - abort_request = partial(self.abort_request, verbose=verbose)

193 - stream = AsyncStream(request_id, abort_request)

194 - self._new_requests.put_nowait((stream, {

195 - "request_id": request_id,

196 - **engine_add_request_kwargs

197 - }))

198 -

199 - self.new_requests_event.set()

200 -

201 - if verbose:

202 - logger.info("Added request %s.", request_id)

203 -

204 - return stream

205 -

206 - def abort_request(self,

207 - request_id: str,

208 - *,

209 - exception: Optional[Union[BaseException,

210 - Type[BaseException]]] = None,

211 - verbose: bool = False) -> None:

212 - """Abort a request during next background loop iteration."""

213 - if verbose:

214 - logger.info("Aborted request %s.", request_id)

215 -

216 - self._aborted_requests.put_nowait(request_id)

217 -

218 - stream = self._request_streams.pop(request_id, None)

219 - if stream is not None:

220 - stream.finish(exception=exception)

221 -

222 - def get_new_and_aborted_requests(self) -> Tuple[List[Dict], Set[str]]:

223 - """Get the new requests and finished requests to be

224 - sent to the engine."""

225 - new_requests: List[Dict] = []

226 - finished_requests: Set[str] = set()

227 -

228 - while not self._aborted_requests.empty():

229 - request_id = self._aborted_requests.get_nowait()

230 - finished_requests.add(request_id)

231 -

232 - while not self._new_requests.empty():

233 - stream, new_request = self._new_requests.get_nowait()

234 - request_id = stream.request_id

235 - if request_id in finished_requests:

236 - # The request has already been aborted.

237 - stream.finish(asyncio.CancelledError)

238 - finished_requests.discard(request_id)

239 - else:

240 - self._request_streams[request_id] = stream

241 - new_requests.append(new_request)

242 -

243 - return new_requests, finished_requests

244 -

245 - async def wait_for_new_requests(self):

246 - if not self.has_new_requests():

247 - await self.new_requests_event.wait()

248 - self.new_requests_event.clear()

249 -

250 - def has_new_requests(self):

251 - return not self._new_requests.empty()

252 -

253 -

254 -class _AsyncLLMEngine(LLMEngine):

255 - """Extension of LLMEngine to add async methods."""

256 -

257 - def __init__(self, *args, **kwargs):

258 - super().__init__(*args, **kwargs)

259 -

260 - async def step_async(self, virtual_engine: int) -> List[RequestOutput]:

261 - """Performs one decoding iteration and returns newly generated results.

262 - The workers are ran asynchronously if possible.

263 -

264 - This function performs one decoding iteration of the engine. It first

265 - schedules the sequences to be executed in the next iteration and the

266 - token blocks to be swapped in/out/copy. Then, it executes the model

267 - and updates the scheduler with the model outputs. Finally, it decodes

268 - the sequences and returns the newly generated results.

269 - """

270 - # these are cached outputs from previous iterations. None if on first

271 - # iteration

272 - cached_outputs = self.cached_scheduler_outputs[virtual_engine]

273 - seq_group_metadata_list = cached_outputs.seq_group_metadata_list

274 - scheduler_outputs = cached_outputs.scheduler_outputs

275 - allow_async_output_proc = cached_outputs.allow_async_output_proc

276 -

277 - ctx = self.scheduler_contexts[virtual_engine]

278 -

279 - # Clear outputs for each new scheduler iteration

280 - ctx.request_outputs.clear()

281 -

282 - # skip the scheduler if there are any remaining steps in the seq groups.

283 - # This ensures that the scheduler is only called again when the current

284 - # batch has completed.

285 - if not self._has_remaining_steps(seq_group_metadata_list):

286 -

287 - # Schedule iteration

288 - (seq_group_metadata_list, scheduler_outputs,

289 - allow_async_output_proc

290 - ) = self.scheduler[virtual_engine].schedule()

291 -

292 - ctx.seq_group_metadata_list = seq_group_metadata_list

293 - ctx.scheduler_outputs = scheduler_outputs

294 -

295 - if not scheduler_outputs.is_empty():

296 - # this will cause mamba_cache/minimax_cache failed

297 - # to release finished_requests_ids of the last steps

298 - finished_requests_ids = self.scheduler[

299 - virtual_engine].get_and_reset_finished_requests_ids()

300 -

301 - # Maybe switch from async mode to sync mode

302 - if not allow_async_output_proc and len(ctx.output_queue) > 0:

303 - self._process_model_outputs(ctx=ctx)

304 -

305 - else:

306 - finished_requests_ids = list()

307 -

308 - assert seq_group_metadata_list is not None

309 - assert scheduler_outputs is not None

310 -

311 - if not scheduler_outputs.is_empty():

312 -

313 - # Check if we have a cached last_output from the previous iteration.

314 - # For supporting PP this is probably the best way to pass the

315 - # sampled_token_ids, as a separate broadcast over all the PP stages

316 - # will cause one virtual engine's microbatch to block the pipeline.

317 - last_sampled_token_ids = \

318 - self._get_last_sampled_token_ids(virtual_engine)

319 -

320 - execute_model_req = ExecuteModelRequest(

321 - seq_group_metadata_list=seq_group_metadata_list,

322 - blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,

323 - blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,

324 - blocks_to_copy=scheduler_outputs.blocks_to_copy,

325 - virtual_engine=virtual_engine,

326 - num_lookahead_slots=scheduler_outputs.num_lookahead_slots,

327 - running_queue_size=scheduler_outputs.running_queue_size,

328 - finished_requests_ids=finished_requests_ids,

329 - # We use ExecuteModelRequest to pass the last sampled_token_ids

330 - # to each of the non-last PP stages for in-place prepare_input.

331 - last_sampled_token_ids=last_sampled_token_ids)

332 -

333 - if allow_async_output_proc:

334 - execute_model_req.async_callback = self.async_callbacks[

335 - virtual_engine]

336 -

337 - # Execute the model.

338 - outputs = await self.model_executor.execute_model_async(

339 - execute_model_req)

340 -

341 - else:

342 - if len(ctx.output_queue) > 0:

343 - self._process_model_outputs(ctx=ctx)

344 - outputs = []

345 -

346 - if not self._has_remaining_steps(seq_group_metadata_list):

347 - # is_first_step_output is True only when the num_steps of all

348 - # the sequences are 1.

349 - is_first_step_output: bool = False if not seq_group_metadata_list \

350 - else seq_group_metadata_list[0].state.num_steps == 1

351 -

352 - ctx.append_output(outputs=outputs,

353 - seq_group_metadata_list=seq_group_metadata_list,

354 - scheduler_outputs=scheduler_outputs,

355 - is_async=allow_async_output_proc,

356 - is_last_step=True,

357 - is_first_step_output=is_first_step_output)

358 -

359 - if outputs and allow_async_output_proc:

360 - assert len(

361 - outputs

362 - ) == 1, "Async postprocessor expects only a single output set"

363 - self._advance_to_next_step(

364 - outputs[0], seq_group_metadata_list,

365 - scheduler_outputs.scheduled_seq_groups)

366 -

367 - if not allow_async_output_proc:

368 - self._process_model_outputs(ctx=ctx)

369 -

370 - # Log stats.

371 - self.do_log_stats(scheduler_outputs, outputs)

372 -

373 - # Tracing

374 - self.do_tracing(scheduler_outputs)

375 -

376 - else:

377 - # Multi-step case

378 - return ctx.request_outputs

379 -

380 - if not self.has_unfinished_requests():

381 - # Drain async postprocessor (if exists)

382 - if len(ctx.output_queue) > 0:

383 - self._process_model_outputs(ctx=ctx)

384 - assert len(ctx.output_queue) == 0

385 -

386 - return ctx.request_outputs

387 -

388 - async def stop_remote_worker_execution_loop_async(self) -> None:

389 - """Stop the remote worker execution loop."""

390 - await self.model_executor.stop_remote_worker_execution_loop_async()

391 -

392 - async def get_tokenizer_async(self) -> AnyTokenizer:

393 - return self.get_tokenizer()

394 -

395 - async def add_request_async(

396 - self,

397 - request_id: str,

398 - prompt: PromptType,

399 - params: SamplingParams,

400 - arrival_time: Optional[float] = None,

401 - lora_request: Optional[LoRARequest] = None,

402 - trace_headers: Optional[Mapping[str, str]] = None,

403 - priority: int = 0,

404 - data_parallel_rank: Optional[int] = None,

405 - tokenization_kwargs: Optional[dict[str, Any]] = None,

406 - ) -> None:

407 - """

408 - Async version of

409 - [`add_request`][vllm.engine.llm_engine.LLMEngine.add_request].

410 - """

411 - if lora_request is not None and not self.lora_config:

412 - raise ValueError(f"Got lora_request {lora_request} but LoRA is "

413 - "not enabled!")

414 - if priority != 0 and not self.scheduler_config.policy == "priority":

415 - raise ValueError(f"Got priority {priority} but "

416 - "Priority scheduling is not enabled.")

417 - if arrival_time is None:

418 - arrival_time = time.time()

419 -

420 - if data_parallel_rank is not None:

421 - raise ValueError("Targeting data_parallel_rank only supported "

422 - "in v1 client.")

423 -

424 - if (isinstance(prompt, dict)

425 - and prompt.get("prompt_embeds", None) is not None

426 - and not prompt.get("prompt_token_ids", None)):

427 - # We use the -2 dimension (instead of 0) in case a batched input

428 - # of batch size 1 is passed in.

429 - prompt["prompt_token_ids"] = [0

430 - ] * prompt["prompt_embeds"].shape[-2]

431 -

432 - processed_inputs = await self.input_preprocessor.preprocess_async(

433 - prompt,

434 - tokenization_kwargs=tokenization_kwargs,

435 - )

436 -

437 - self._add_processed_request(

438 - request_id=request_id,

439 - processed_inputs=processed_inputs,

440 - params=params,

441 - arrival_time=arrival_time,

442 - lora_request=lora_request,

443 - trace_headers=trace_headers,

444 - priority=priority,

445 - )

446 -

447 - async def check_health_async(self) -> None:

448 - self.model_executor.check_health()

449 -

450 - async def collective_rpc_async(self,

451 - method: str,

452 - timeout: Optional[float] = None,

453 - args: tuple = (),

454 - kwargs: Optional[dict] = None):

455 - raise NotImplementedError

456 -

457 -

458 -class AsyncLLMEngine(EngineClient):

459 - """An asynchronous wrapper for [`LLMEngine`][vllm.LLMEngine].

460 -

461 - This class is used to wrap the [`LLMEngine`][vllm.LLMEngine] class to

462 - make it asynchronous. It uses asyncio to create a background loop that keeps

463 - processing incoming requests. The [`LLMEngine`][vllm.LLMEngine] is kicked

464 - by the generate method when there are requests in the waiting queue. The

465 - generate method yields the outputs from the [`LLMEngine`][vllm.LLMEngine]

466 - to the caller.

467 -

468 - Args:

469 - log_requests: Whether to log the requests.

470 - start_engine_loop: If True, the background task to run the engine

471 - will be automatically started in the generate call.

472 - *args: Arguments for [`LLMEngine`][vllm.LLMEngine].

473 - **kwargs: Arguments for [`LLMEngine`][vllm.LLMEngine].

474 - """

475 -

476 - _engine_class: Type[_AsyncLLMEngine] = _AsyncLLMEngine

477 -

478 - def __init__(self,

479 - *args: Any,

480 - log_requests: bool = True,

481 - start_engine_loop: bool = True,

482 - **kwargs: Any) -> None:

483 - if envs.VLLM_USE_V1:

484 - raise ValueError(

485 - "Using V0 AsyncLLMEngine, but envs.VLLM_USE_V1=True. "

486 - "This should not happen. As a workaround, try using "

487 - "AsyncLLMEngine.from_vllm_config(...) or explicitly set "

488 - "VLLM_USE_V1=0 or 1 and report this issue on Github.")

489 -

490 - self.log_requests = log_requests

491 - self.engine = self._engine_class(*args, **kwargs)

492 -

493 - # This ensures quick processing of request outputs

494 - # so the append to asyncio queues is not delayed,

495 - # especially for multi-step.

496 - self.use_process_request_outputs_callback = (

497 - self.engine.model_config.use_async_output_proc)

498 -

499 - if self.use_process_request_outputs_callback:

500 - self.engine.process_request_outputs_callback = \

501 - weak_bind(self.process_request_outputs)

502 -

503 - self.background_loop: Optional[asyncio.Future] = None

504 - # We need to keep a reference to unshielded

505 - # task as well to prevent it from being garbage

506 - # collected

507 - self._background_loop_unshielded: Optional[asyncio.Task] = None

508 - self.start_engine_loop = start_engine_loop

509 - self._errored_with: Optional[BaseException] = None

510 -

511 - # Lazy initialized fields

512 - self._request_tracker: RequestTracker

513 -

514 - def __del__(self):

515 - if rt := getattr(self, "request_tracker", None):

516 - # Wake up engine loop so that it will exit cleanly

517 - rt.new_requests_event.set()

518 -

519 - @classmethod

520 - def _get_executor_cls(cls,

521 - engine_config: VllmConfig) -> Type[ExecutorBase]:

522 - return LLMEngine._get_executor_cls(engine_config)

523 -

524 - @classmethod

525 - @deprecate_kwargs(

526 - "disable_log_requests",

527 - additional_message=("This argument will have no effect. "

528 - "Use `enable_log_requests` instead."),

529 - )

530 - def from_vllm_config(

531 - cls,

532 - vllm_config: VllmConfig,

533 - start_engine_loop: bool = True,

534 - usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

535 - stat_loggers: Optional[dict[str, StatLoggerBase]] = None,

536 - enable_log_requests: bool = False,

537 - disable_log_stats: bool = False,

538 - disable_log_requests: bool = True, # Deprecated, will be removed

539 - ) -> "AsyncLLMEngine":

540 - """Create an AsyncLLMEngine from the EngineArgs."""

541 -

542 - return cls(

543 - vllm_config=vllm_config,

544 - executor_class=cls._get_executor_cls(vllm_config),

545 - start_engine_loop=start_engine_loop,

546 - log_requests=enable_log_requests,

547 - log_stats=not disable_log_stats,

548 - usage_context=usage_context,

549 - stat_loggers=stat_loggers,

550 - )

551 -

552 - @classmethod

553 - def from_engine_args(

554 - cls,

555 - engine_args: AsyncEngineArgs,

556 - start_engine_loop: bool = True,

557 - usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,

558 - stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,

559 - ) -> "AsyncLLMEngine":

560 - """Creates an async LLM engine from the engine arguments."""

561 -

562 - vllm_config = engine_args.create_engine_config(usage_context)

563 -

564 - async_engine_cls = cls

565 - if envs.VLLM_USE_V1:

566 - from vllm.v1.engine.async_llm import AsyncLLM as V1AsyncLLMEngine

567 - async_engine_cls = V1AsyncLLMEngine

568 -

569 - return async_engine_cls.from_vllm_config(

570 - vllm_config=vllm_config,

571 - start_engine_loop=start_engine_loop,

572 - usage_context=usage_context,

573 - stat_loggers=stat_loggers,

574 - disable_log_stats=engine_args.disable_log_stats,

575 - enable_log_requests=engine_args.enable_log_requests,

576 - )

577 -

578 - @property

579 - def is_running(self) -> bool:

580 - return (self.background_loop is not None

581 - and self._background_loop_unshielded is not None

582 - and not self._background_loop_unshielded.done())

583 -

584 - @property

585 - def is_stopped(self) -> bool:

586 - return self.errored or (self.background_loop is not None and

587 - self._background_loop_unshielded is not None

588 - and self._background_loop_unshielded.done())

589 -

590 - @property

591 - def errored(self) -> bool:

592 - return self._errored_with is not None

593 -

594 - @property

595 - def dead_error(self) -> BaseException:

596 - return AsyncEngineDeadError(

597 - "Background loop is not running. If it was running, "

598 - "inspect the output to find the stacktrace of the "

599 - "error that caused the background loop to stop "

600 - "(AsyncEngineDeadError).")

601 -

602 - def set_errored(self, exc: Exception) -> None:

603 - self._errored_with = exc

604 -

605 - def _error_callback(self, exc: Exception) -> None:

606 - self.set_errored(exc)

607 - self._request_tracker.propagate_exception(exc)

608 -

609 - async def get_input_preprocessor(self) -> InputPreprocessor:

610 - return self.engine.input_preprocessor

611 -

612 - async def get_tokenizer(self) -> AnyTokenizer:

613 - return self.engine.get_tokenizer()

614 -

615 - def start_background_loop(self) -> None:

616 - """Start the background loop."""

617 - if self.errored:

618 - raise AsyncEngineDeadError(

619 - "Background loop has errored already.") from self._errored_with

620 - if self.is_running:

621 - raise RuntimeError("Background loop is already running.")

622 - # Initialize the RequestTracker here so it uses the right event loop.

623 - self._request_tracker = RequestTracker()

624 -

625 - self._background_loop_unshielded = asyncio.get_event_loop(

626 - ).create_task(self.run_engine_loop(weakref.ref(self)))

627 - self._background_loop_unshielded.add_done_callback(

628 - partial(_log_task_completion, error_callback=self._error_callback))

629 - self.background_loop = asyncio.shield(self._background_loop_unshielded)

630 -

631 - def shutdown_background_loop(self) -> None:

632 - """

633 - Shut down the background loop.

634 -

635 - This method needs to be called during cleanup to remove

636 - references to `self` and properly GC the resources held

637 - by the async LLM engine (e.g., the executors as well as

638 - their resources).

639 - """

640 - if self._background_loop_unshielded is not None:

641 - self._background_loop_unshielded.cancel()

642 - self._background_loop_unshielded = None

643 - self.background_loop = None

644 -

645 - async def engine_step(self, virtual_engine: int) -> bool:

646 - """Kick the engine to process the waiting requests.

647 -

648 - Returns True if there are in-progress requests."""

649 -

650 - new_requests, aborted_requests = (

651 - self._request_tracker.get_new_and_aborted_requests())

652 -

653 - for new_request in new_requests:

654 - # Add the request into the vLLM engine's waiting queue.

655 - try:

656 - await self.engine.add_request_async(**new_request)

657 - except ValueError as e:

658 - # TODO: use a vLLM specific error for failed validation

659 - self._request_tracker.process_exception(

660 - new_request["request_id"],

661 - e,

662 - verbose=self.log_requests,

663 - )

664 -

665 - if aborted_requests:

666 - await self._engine_abort(aborted_requests)

667 -

668 - request_outputs = await self.engine.step_async(virtual_engine)

669 -

670 - # Put the outputs into the corresponding streams.

671 - # If used as a callback, then already invoked inside

672 - # LLMEngine's _process_model_outputs

673 - if not self.use_process_request_outputs_callback:

674 - all_finished = self.process_request_outputs(request_outputs)

675 - else:

676 - # For callback case, we only need to detect when all

677 - # requests are finished

678 - all_finished = all(request_output.finished

679 - for request_output in request_outputs)

680 -

681 - return not all_finished

682 -

683 - def process_request_outputs(self, request_outputs) -> bool:

684 - # Put the outputs into the corresponding streams.

685 - all_finished = True

686 - for request_output in request_outputs:

687 - self._request_tracker.process_request_output(

688 - request_output, verbose=self.log_requests)

689 - all_finished = all_finished and request_output.finished

690 -

691 - return all_finished

692 -

693 - async def _engine_abort(self, request_ids: Iterable[str]):

694 - self.engine.abort_request(request_ids)

695 -

696 - @staticmethod

697 - async def run_engine_loop(engine_ref: ReferenceType):

698 - """We use a weakref to the engine so that the running loop

699 - doesn't prevent the engine being garbage collected."""

700 - engine: Optional[AsyncLLMEngine] = engine_ref()

701 - if not engine:

702 - return

703 -

704 - pipeline_parallel_size = \

705 - engine.engine.parallel_config.pipeline_parallel_size

706 - has_requests_in_progress = [False] * pipeline_parallel_size

707 - while True:

708 - if not any(has_requests_in_progress):

709 - logger.debug("Waiting for new requests...")

710 - # Stop the execute model loop in parallel workers until there

711 - # are more requests to process. This avoids waiting

712 - # indefinitely in torch.distributed ops which may otherwise

713 - # time out, and unblocks the RPC thread in the workers so that

714 - # they can process any other queued control plane messages,

715 - # such as add/remove lora adapters.

716 - await engine.engine.stop_remote_worker_execution_loop_async()

717 - request_tracker = engine._request_tracker

718 - # Allow engine to be garbage collected while

719 - # waiting for new requests

720 - del engine

721 - await asyncio.sleep(0)

722 - if engine_ref() is None:

723 - return

724 - await request_tracker.wait_for_new_requests()

725 - engine = engine_ref()

726 - if not engine:

727 - return

728 - logger.debug("Got new requests!")

729 - requests_in_progress = [

730 - asyncio.create_task(engine.engine_step(ve))

731 - for ve in range(pipeline_parallel_size)

732 - ]

733 - has_requests_in_progress = [True] * pipeline_parallel_size

734 -

735 - # Abort if iteration takes too long due to unrecoverable errors

736 - # (eg. NCCL timeouts).

737 - try:

738 - async with asyncio_timeout(ENGINE_ITERATION_TIMEOUT_S):

739 - done, _ = await asyncio.wait(

740 - requests_in_progress,

741 - return_when=asyncio.FIRST_COMPLETED)

742 - for _ in range(pipeline_parallel_size):

743 - await asyncio.sleep(0)

744 - for task in done:

745 - result = task.result()

746 - virtual_engine = requests_in_progress.index(task)

747 - has_unfinished_requests = (

748 - engine.engine.

749 - has_unfinished_requests_for_virtual_engine(

750 - virtual_engine))

751 - if result or has_unfinished_requests:

752 - requests_in_progress[virtual_engine] = (

753 - asyncio.create_task(

754 - engine.engine_step(virtual_engine)))

755 - has_requests_in_progress[virtual_engine] = True

756 - else:

757 - has_requests_in_progress[virtual_engine] = False

758 - except asyncio.TimeoutError as exc:

759 - logger.error(

760 - "Engine iteration timed out. This should never happen!")

761 - engine.set_errored(exc)

762 - raise

763 - await asyncio.sleep(0)

764 -

765 - async def add_request(

766 - self,

767 - request_id: str,

768 - prompt: PromptType,

769 - params: SamplingParams,

770 - arrival_time: Optional[float] = None,

771 - lora_request: Optional[LoRARequest] = None,

772 - trace_headers: Optional[Mapping[str, str]] = None,

773 - priority: int = 0,

774 - data_parallel_rank: Optional[int] = None,

775 - tokenization_kwargs: Optional[dict[str, Any]] = None,

776 - ) -> AsyncGenerator[RequestOutput, None]:

777 - if not self.is_running:

778 - if self.start_engine_loop:

779 - self.start_background_loop()

780 - else:

781 - raise AsyncEngineDeadError(

782 - "Background loop is not running. If it was running, "

783 - "inspect the output to find the stacktrace of the "

784 - "error that caused the background loop to stop "

785 - "(AsyncEngineDeadError).")

786 -

787 - if (priority != 0

788 - and not self.engine.scheduler_config.policy == "priority"):

789 - raise ValueError(f"Got priority {priority} but "

790 - "Priority scheduling is not enabled.")

791 -

792 - stream = self._request_tracker.add_request(

793 - request_id,

794 - verbose=self.log_requests,

795 - prompt=prompt,

796 - params=params,

797 - arrival_time=arrival_time or time.time(),

798 - lora_request=lora_request,

799 - trace_headers=trace_headers,

800 - priority=priority,

801 - data_parallel_rank=data_parallel_rank,

802 - tokenization_kwargs=tokenization_kwargs,

803 - )

804 -

805 - return stream.generator()

806 -

807 - async def generate(

808 - self,

809 - prompt: PromptType,

810 - sampling_params: SamplingParams,

811 - request_id: str,

812 - lora_request: Optional[LoRARequest] = None,

813 - trace_headers: Optional[Mapping[str, str]] = None,

814 - priority: int = 0,

815 - data_parallel_rank: Optional[int] = None,

816 - ) -> AsyncGenerator[RequestOutput, None]:

817 - """Generate outputs for a request.

818 -

819 - Generate outputs for a request. This method is a coroutine. It adds the

820 - request into the waiting queue of the LLMEngine and streams the outputs

821 - from the LLMEngine to the caller.

822 -

823 - Args:

824 - prompt: The prompt to the LLM. See

825 - [`PromptType`][vllm.inputs.PromptType] for more details about

826 - the format of each input.

827 - sampling_params: The sampling parameters of the request.

828 - request_id: The unique id of the request.

829 - lora_request: LoRA request to use for generation, if any.

830 - trace_headers: OpenTelemetry trace headers.

831 - priority: The priority of the request.

832 - Only applicable with priority scheduling.

833 - data_parallel_rank: The (global) data parallel rank that must

834 - handle this request. Only applicable if DP is enabled.

835 - Yields:

836 - The output `RequestOutput` objects from the LLMEngine

837 - for the request.

838 -

839 - Details:

840 - - If the engine is not running, start the background loop,

841 - which iteratively invokes

842 - [`engine_step`][vllm.engine.async_llm_engine.AsyncLLMEngine.engine_step]

843 - to process the waiting requests.

844 - - Add the request to the engine's `RequestTracker`.

845 - On the next background loop, this request will be sent to

846 - the underlying engine.

847 - Also, a corresponding `AsyncStream` will be created.

848 - - Wait for the request outputs from `AsyncStream` and yield them.

849 -

850 - Example:

851 - >>> # Please refer to entrypoints/api_server.py for

852 - >>> # the complete example.

853 - >>>

854 - >>> # initialize the engine and the example input

855 - >>> # note that engine_args here is AsyncEngineArgs instance

856 - >>> engine = AsyncLLMEngine.from_engine_args(engine_args)

857 - >>> example_input = {

858 - >>> "prompt": "What is LLM?",

859 - >>> "stream": False, # assume the non-streaming case

860 - >>> "temperature": 0.0,

861 - >>> "request_id": 0,

862 - >>> }

863 - >>>

864 - >>> # start the generation

865 - >>> results_generator = engine.generate(

866 - >>> example_input["prompt"],

867 - >>> SamplingParams(temperature=example_input["temperature"]),

868 - >>> example_input["request_id"])

869 - >>>

870 - >>> # get the results

871 - >>> final_output = None

872 - >>> async for request_output in results_generator:

873 - >>> if await request.is_disconnected():

874 - >>> # Abort the request if the client disconnects.

875 - >>> await engine.abort(request_id)

876 - >>> # Return or raise an error

877 - >>> ...

878 - >>> final_output = request_output

879 - >>>

880 - >>> # Process and return the final output

881 - >>> ...

882 - """

883 - try:

884 - async for output in await self.add_request(

885 - request_id,

886 - prompt,

887 - sampling_params,

888 - lora_request=lora_request,

889 - trace_headers=trace_headers,

890 - priority=priority,

891 - data_parallel_rank=data_parallel_rank,

892 - ):

893 - yield LLMEngine.validate_output(output, RequestOutput)

894 - except asyncio.CancelledError:

895 - await self.abort(request_id)

896 - raise

897 -

898 - def encode(

899 - self,

900 - prompt: PromptType,

901 - pooling_params: PoolingParams,

902 - request_id: str,

903 - lora_request: Optional[LoRARequest] = None,

904 - trace_headers: Optional[Mapping[str, str]] = None,

905 - priority: int = 0,

906 - tokenization_kwargs: Optional[dict[str, Any]] = None,

907 - ) -> AsyncGenerator[PoolingRequestOutput, None]:

908 - raise NotImplementedError(

909 - "Pooling models are not supported in vLLM V0")

910 -

911 - async def abort(self, request_id: Union[str, Iterable[str]]) -> None:

912 - """Abort a request.

913 -

914 - Abort a submitted request. If the request is finished or not found,

915 - this method will be a no-op.

916 -

917 - Args:

918 - request_id: The unique id of the request.

919 - """

920 - if not isinstance(request_id, str):

921 - raise RuntimeError("Only single-request abort supported in"

922 - " deprecated V0")

923 - if not self.is_running:

924 - raise AsyncEngineDeadError(

925 - "Background loop is not running. If it was running, "

926 - "inspect the output to find the stacktrace of the "

927 - "error that caused the background loop to stop "

928 - "(AsyncEngineDeadError).")

929 -

930 - return self._abort(request_id)

931 -

932 - def _abort(self, request_id: str) -> None:

933 - """Abort a request.

934 -

935 - Abort a submitted request. If the request is finished or not found,

936 - this method will be a no-op.

937 -

938 - Args:

939 - request_id: The unique id of the request.

940 - """

941 - self._request_tracker.abort_request(request_id,

942 - exception=asyncio.CancelledError,

943 - verbose=self.log_requests)

944 -

945 - async def get_vllm_config(self) -> VllmConfig:

946 - """Get the vllm configuration of the vLLM engine."""

947 - return self.engine.get_vllm_config()

948 -

949 - async def get_model_config(self) -> ModelConfig:

950 - """Get the model configuration of the vLLM engine."""

951 - return self.engine.get_model_config()

952 -

953 - async def get_parallel_config(self) -> ParallelConfig:

954 - """Get the parallel configuration of the vLLM engine."""

955 - return self.engine.get_parallel_config()

956 -

957 - async def get_scheduler_config(self) -> SchedulerConfig:

958 - """Get the scheduling configuration of the vLLM engine."""

959 - return self.engine.get_scheduler_config()

960 -

961 - async def get_lora_config(self) -> LoRAConfig:

962 - """Get the lora configuration of the vLLM engine."""

963 - return self.engine.get_lora_config()

964 -

965 - async def do_log_stats(

966 - self,

967 - scheduler_outputs: Optional[SchedulerOutputs] = None,

968 - model_output: Optional[List[SamplerOutput]] = None) -> None:

969 - self.engine.do_log_stats()

970 -

971 - async def check_health(self) -> None:

972 - """Raises an error if engine is unhealthy."""

973 - t = time.perf_counter()

974 - logger.debug("Starting health check...")

975 - if self.is_stopped:

976 - raise AsyncEngineDeadError("Background loop is stopped.")

977 -

978 - await self.engine.check_health_async()

979 - logger.debug("Health check took %fs", time.perf_counter() - t)

980 -

981 - async def is_tracing_enabled(self) -> bool:

982 - return self.engine.is_tracing_enabled()

983 -

984 - def add_logger(self, logger_name: str, logger: StatLoggerBase) -> None:

985 - self.engine.add_logger(logger_name=logger_name, logger=logger)

986 -

987 - def remove_logger(self, logger_name: str) -> None:

988 - self.engine.remove_logger(logger_name=logger_name)

989 -

990 - async def start_profile(self) -> None:

991 - self.engine.start_profile()

992 -

993 - async def stop_profile(self) -> None:

994 - self.engine.stop_profile()

995 -

996 - async def reset_mm_cache(self) -> None:

997 - self.engine.reset_mm_cache()

998 -

999 - async def reset_prefix_cache(self,

1000 - device: Optional[Device] = None) -> None:

1001 - self.engine.reset_prefix_cache(device)

1002 -

1003 - async def sleep(self, level: int = 1) -> None:

1004 - await self.reset_prefix_cache()

1005 - self.engine.sleep(level)

1006 -

1007 - async def wake_up(self, tags: Optional[list[str]] = None) -> None:

1008 - self.engine.wake_up(tags)

1009 -

1010 - async def is_sleeping(self) -> bool:

1011 - return self.engine.is_sleeping()

1012 -

1013 - async def add_lora(self, lora_request: LoRARequest) -> bool:

1014 - return self.engine.add_lora(lora_request)

1015 -

1016 - async def collective_rpc(self,

1017 - method: str,

1018 - timeout: Optional[float] = None,

1019 - args: tuple = (),

1020 - kwargs: Optional[dict] = None):

1021 - """

1022 - Perform a collective RPC call to the given path.

1023 - """

1024 - return await self.engine.collective_rpc_async(method, timeout, args,

1025 - kwargs)

1026 -

1027 -

1028 -# TODO(v1): Remove this class proxy when V1 goes default.

1029 -if envs.is_set("VLLM_USE_V1") and envs.VLLM_USE_V1:

1030 - from vllm.v1.engine.async_llm import AsyncLLM

1031 -

1032 - AsyncLLMEngine = AsyncLLM # type: ignore

6 +AsyncLLMEngine = AsyncLLM # type: ignore

vllm/entrypoints/launcher.py

-2

@@ -11,7 +11,6 @@ import uvicorn

11 from fastapi import FastAPI, Request, Response

13 from vllm import envs

14 -from vllm.engine.async_llm_engine import AsyncEngineDeadError

14 from vllm.engine.protocol import EngineClient

15 from vllm.entrypoints.constants import (H11_MAX_HEADER_COUNT_DEFAULT,

16 H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT)

@@ -154,7 +153,6 @@ def _add_shutdown_handlers(app: FastAPI, server: uvicorn.Server) -> None:

153 """

154

155 @app.exception_handler(RuntimeError)

157 - @app.exception_handler(AsyncEngineDeadError)

156 @app.exception_handler(EngineDeadError)

157 @app.exception_handler(EngineGenerateError)

158 async def runtime_exception_handler(request: Request, __):

vllm/entrypoints/openai/api_server.py

+26 -43

@@ -38,7 +38,6 @@ from typing_extensions import assert_never

38 import vllm.envs as envs

39 from vllm.config import VllmConfig

40 from vllm.engine.arg_utils import AsyncEngineArgs

41 -from vllm.engine.async_llm_engine import AsyncLLMEngine # type: ignore

41 from vllm.engine.protocol import EngineClient

42 from vllm.entrypoints.chat_utils import (load_chat_template,

43 resolve_hf_chat_template,

@@ -201,50 +200,34 @@ async def build_async_engine_client_from_engine_args(

200 vllm_config = engine_args.create_engine_config(usage_context=usage_context)

201

202 # V1 AsyncLLM.

204 - if envs.VLLM_USE_V1:

205 - if disable_frontend_multiprocessing:

206 - logger.warning(

207 - "V1 is enabled, but got --disable-frontend-multiprocessing. "

208 - "To disable frontend multiprocessing, set VLLM_USE_V1=0.")

209 -

210 - from vllm.v1.engine.async_llm import AsyncLLM

211 - async_llm: Optional[AsyncLLM] = None

212 - client_count = client_config.pop(

213 - "client_count") if client_config else 1

214 - client_index = client_config.pop(

215 - "client_index") if client_config else 0

216 - try:

217 - async_llm = AsyncLLM.from_vllm_config(

218 - vllm_config=vllm_config,

219 - usage_context=usage_context,

220 - enable_log_requests=engine_args.enable_log_requests,

221 - disable_log_stats=engine_args.disable_log_stats,

222 - client_addresses=client_config,

223 - client_count=client_count,

224 - client_index=client_index)

225 -

226 - # Don't keep the dummy data in memory

227 - await async_llm.reset_mm_cache()

228 -

229 - yield async_llm

230 - finally:

231 - if async_llm:

232 - async_llm.shutdown()

203 + assert envs.VLLM_USE_V1

204

234 - # V0 AsyncLLM.

235 - else:

205 + if disable_frontend_multiprocessing:

206 + logger.warning(

207 + "V1 is enabled, but got --disable-frontend-multiprocessing. "

208 + "To disable frontend multiprocessing, set VLLM_USE_V1=0.")

209

237 - engine_client: Optional[EngineClient] = None

238 - try:

239 - engine_client = AsyncLLMEngine.from_vllm_config(

240 - vllm_config=vllm_config,

241 - usage_context=usage_context,

242 - enable_log_requests=engine_args.enable_log_requests,

243 - disable_log_stats=engine_args.disable_log_stats)

244 - yield engine_client

245 - finally:

246 - if engine_client and hasattr(engine_client, "shutdown"):

247 - engine_client.shutdown()

210 + from vllm.v1.engine.async_llm import AsyncLLM

211 + async_llm: Optional[AsyncLLM] = None

212 + client_count = client_config.pop("client_count") if client_config else 1

213 + client_index = client_config.pop("client_index") if client_config else 0

214 + try:

215 + async_llm = AsyncLLM.from_vllm_config(

216 + vllm_config=vllm_config,

217 + usage_context=usage_context,

218 + enable_log_requests=engine_args.enable_log_requests,

219 + disable_log_stats=engine_args.disable_log_stats,

220 + client_addresses=client_config,

221 + client_count=client_count,

222 + client_index=client_index)

223 +

224 + # Don't keep the dummy data in memory

225 + await async_llm.reset_mm_cache()

226 +

227 + yield async_llm

228 + finally:

229 + if async_llm:

230 + async_llm.shutdown()

231

232

233 async def validate_json_request(raw_request: Request):