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

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 MQLLMEngine (#25019)

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

Woosuk Kwon committed Sep 16, 2025 at 21:29 UTC 5801e4977679895d89493e84088e7936f528285f

15 files changed +12 -1969

.buildkite/test-pipeline.yaml

-2

@@ -46,7 +46,6 @@ steps:

46 mirror_hardwares: [amdexperimental]

47 source_file_dependencies:

48 - vllm/

49 - - tests/mq_llm_engine

49 - tests/async_engine

50 - tests/test_inputs.py

51 - tests/test_outputs.py

@@ -57,7 +56,6 @@ steps:

56 - tests/transformers_utils

57 commands:

58 - python3 standalone_tests/lazy_imports.py

60 - - pytest -v -s mq_llm_engine # MQLLMEngine

59 - pytest -v -s async_engine # AsyncLLMEngine

60 - pytest -v -s test_inputs.py

61 - pytest -v -s test_outputs.py

tests/entrypoints/openai/test_lora_resolvers.py

+2 -2

@@ -10,7 +10,6 @@ from unittest.mock import MagicMock

10 import pytest

12 from vllm.config.multimodal import MultiModalConfig

13 -from vllm.engine.multiprocessing.client import MQLLMEngineClient

13 from vllm.entrypoints.openai.protocol import CompletionRequest, ErrorResponse

14 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion

15 from vllm.entrypoints.openai.serving_models import (BaseModelPath,

@@ -18,6 +17,7 @@ from vllm.entrypoints.openai.serving_models import (BaseModelPath,

17 from vllm.lora.request import LoRARequest

18 from vllm.lora.resolver import LoRAResolver, LoRAResolverRegistry

19 from vllm.transformers_utils.tokenizer import get_tokenizer

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

22 MODEL_NAME = "openai-community/gpt2"

23 BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]

@@ -82,7 +82,7 @@ def register_mock_resolver():

82 @pytest.fixture

83 def mock_serving_setup():

84 """Provides a mocked engine and serving completion instance."""

85 - mock_engine = MagicMock(spec=MQLLMEngineClient)

85 + mock_engine = MagicMock(spec=AsyncLLM)

86 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

87 mock_engine.errored = False

tests/entrypoints/openai/test_serving_chat.py

+7 -7

@@ -13,12 +13,12 @@ import pytest

13 import pytest_asyncio

15 from vllm.config.multimodal import MultiModalConfig

16 -from vllm.engine.multiprocessing.client import MQLLMEngineClient

16 from vllm.entrypoints.openai.protocol import ChatCompletionRequest

17 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat

18 from vllm.entrypoints.openai.serving_models import (BaseModelPath,

19 OpenAIServingModels)

20 from vllm.transformers_utils.tokenizer import get_tokenizer

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

23 from ...utils import RemoteOpenAIServer

@@ -276,7 +276,7 @@ def test_async_serving_chat_init():

276

277 @pytest.mark.asyncio

278 async def test_serving_chat_returns_correct_model_name():

279 - mock_engine = MagicMock(spec=MQLLMEngineClient)

279 + mock_engine = MagicMock(spec=AsyncLLM)

280 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

281 mock_engine.errored = False

282

@@ -312,7 +312,7 @@ async def test_serving_chat_returns_correct_model_name():

312

313 @pytest.mark.asyncio

314 async def test_serving_chat_should_set_correct_max_tokens():

315 - mock_engine = MagicMock(spec=MQLLMEngineClient)

315 + mock_engine = MagicMock(spec=AsyncLLM)

316 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

317 mock_engine.errored = False

318

@@ -355,7 +355,7 @@ async def test_serving_chat_should_set_correct_max_tokens():

355 }

356

357 # Reinitialize the engine with new settings

358 - mock_engine = MagicMock(spec=MQLLMEngineClient)

358 + mock_engine = MagicMock(spec=AsyncLLM)

359 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

360 mock_engine.errored = False

361

@@ -410,7 +410,7 @@ async def test_serving_chat_should_set_correct_max_tokens():

410 }

411

412 # Reinitialize the engine with new settings

413 - mock_engine = MagicMock(spec=MQLLMEngineClient)

413 + mock_engine = MagicMock(spec=AsyncLLM)

414 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

415 mock_engine.errored = False

416

@@ -467,7 +467,7 @@ async def test_serving_chat_could_load_correct_generation_config():

467 "repetition_penalty": 1.05

468 }

469

470 - mock_engine = MagicMock(spec=MQLLMEngineClient)

470 + mock_engine = MagicMock(spec=AsyncLLM)

471 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

472 mock_engine.errored = False

473

@@ -523,7 +523,7 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type):

523 mock_model_config = MockModelConfig()

524 mock_model_config.hf_config.model_type = model_type

525

526 - mock_engine = MagicMock(spec=MQLLMEngineClient)

526 + mock_engine = MagicMock(spec=AsyncLLM)

527 mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

528 mock_engine.errored = False

529

tests/mq_llm_engine/__init__.py

tests/mq_llm_engine/conftest.py deleted

-12

@@ -1,12 +0,0 @@

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

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

3 -import pytest

4 -

5 -

6 -@pytest.fixture(scope="function", autouse=True)

7 -def use_v0_only(monkeypatch):

8 - """

9 - Since this module is V0 only, set VLLM_USE_V1=0 for

10 - all tests in the module.

11 - """

12 - monkeypatch.setenv('VLLM_USE_V1', '0')

tests/mq_llm_engine/test_abort.py deleted

-69

@@ -1,69 +0,0 @@

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

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

3 -"""Test that aborting is handled properly."""

4 -

5 -import asyncio

6 -import tempfile

7 -import uuid

8 -

9 -import pytest

10 -

11 -from tests.mq_llm_engine.utils import RemoteMQLLMEngine, generate

12 -from vllm.engine.arg_utils import AsyncEngineArgs

13 -

14 -MODEL = "google/gemma-1.1-2b-it"

15 -ENGINE_ARGS = AsyncEngineArgs(model=MODEL)

16 -RAISED_ERROR = KeyError

17 -RAISED_VALUE = "foo"

18 -EXPECTED_TOKENS = 250

19 -

20 -

21 -@pytest.fixture(scope="function")

22 -def tmp_socket():

23 - with tempfile.TemporaryDirectory() as td:

24 - yield f"ipc://{td}/{uuid.uuid4()}"

25 -

26 -

27 -@pytest.mark.asyncio

28 -async def test_abort(tmp_socket):

29 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

30 - ipc_path=tmp_socket) as engine:

31 -

32 - client = await engine.make_client()

33 -

34 - request_id_to_be_aborted = "request-aborted"

35 - request_ids_a = [f"request-a-{idx}" for idx in range(10)]

36 - request_ids_b = [f"request-b-{idx}" for idx in range(10)]

37 -

38 - # Requests started before one to be aborted.

39 - tasks = []

40 - for request_id in request_ids_a:

41 - tasks.append(

42 - asyncio.create_task(

43 - generate(client, request_id, EXPECTED_TOKENS)))

44 -

45 - # Aborted.

46 - task_aborted = asyncio.create_task(

47 - generate(client, request_id_to_be_aborted, EXPECTED_TOKENS))

48 -

49 - # Requests started after one to be aborted.

50 - for request_id in request_ids_b:

51 - tasks.append(

52 - asyncio.create_task(

53 - generate(client, request_id, EXPECTED_TOKENS)))

54 -

55 - # Actually abort.

56 - await asyncio.sleep(0.5)

57 - await client.abort(request_id_to_be_aborted)

58 -

59 - # Confirm that we got all the EXPECTED tokens from the requests.

60 - for task in tasks:

61 - count, request_id = await task

62 - assert count == EXPECTED_TOKENS, (

63 - f"{request_id} generated only {count} tokens")

64 -

65 - # Cancel task (this will hang indefinitely if not).

66 - task_aborted.cancel()

67 -

68 - # Shutdown.

69 - client.close()

tests/mq_llm_engine/test_error_handling.py deleted

-376

@@ -1,376 +0,0 @@

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

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

3 -"""Test that various errors are handled properly."""

4 -

5 -import asyncio

6 -import tempfile

7 -import time

8 -import uuid

9 -from unittest.mock import Mock

10 -

11 -import pytest

12 -

13 -from tests.mq_llm_engine.utils import RemoteMQLLMEngine

14 -from vllm import SamplingParams

15 -from vllm.engine.arg_utils import AsyncEngineArgs

16 -from vllm.engine.llm_engine import LLMEngine

17 -from vllm.engine.multiprocessing import MQEngineDeadError

18 -from vllm.engine.multiprocessing.engine import MQLLMEngine

19 -from vllm.entrypoints.openai.api_server import build_async_engine_client

20 -from vllm.entrypoints.openai.cli_args import make_arg_parser

21 -from vllm.lora.request import LoRARequest

22 -from vllm.sequence import SequenceGroupMetadata

23 -from vllm.usage.usage_lib import UsageContext

24 -from vllm.utils import FlexibleArgumentParser

25 -

26 -MODEL = "google/gemma-1.1-2b-it"

27 -ENGINE_ARGS = AsyncEngineArgs(model=MODEL, enforce_eager=True)

28 -RAISED_ERROR = KeyError

29 -RAISED_VALUE = "foo"

30 -

31 -

32 -@pytest.fixture(scope="function")

33 -def tmp_socket():

34 - with tempfile.TemporaryDirectory() as td:

35 - yield f"ipc://{td}/{uuid.uuid4()}"

36 -

37 -

38 -def run_with_evil_forward(engine_args: AsyncEngineArgs, ipc_path: str):

39 - # Make engine.

40 - engine = MQLLMEngine.from_engine_args(

41 - engine_args=engine_args,

42 - usage_context=UsageContext.UNKNOWN_CONTEXT,

43 - ipc_path=ipc_path)

44 -

45 - # Raise error during first forward pass.

46 - engine.engine.model_executor.execute_model = Mock(

47 - side_effect=RAISED_ERROR(RAISED_VALUE))

48 -

49 - # Run engine.

50 - engine.start()

51 -

52 -

53 -@pytest.mark.asyncio

54 -async def test_evil_forward(tmp_socket):

55 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

56 - ipc_path=tmp_socket,

57 - run_fn=run_with_evil_forward) as engine:

58 -

59 - client = await engine.make_client()

60 -

61 - # Server should be healthy after initial probe.

62 - await asyncio.sleep(2.0)

63 - await client.check_health()

64 -

65 - # Throws an error that should get ENGINE_DEAD_ERROR.

66 - with pytest.raises(MQEngineDeadError):

67 - async for _ in client.generate(prompt="Hello my name is",

68 - sampling_params=SamplingParams(),

69 - request_id=str(uuid.uuid4())):

70 - pass

71 - assert client.errored

72 -

73 - await asyncio.sleep(1.0)

74 - with pytest.raises(RAISED_ERROR):

75 - await client.check_health()

76 - assert client.errored

77 -

78 - # Shutdown.

79 - client.close()

80 -

81 -

82 -def run_with_evil_model_executor_health(engine_args: AsyncEngineArgs,

83 - ipc_path: str):

84 - # Make engine.

85 - engine = MQLLMEngine.from_engine_args(

86 - engine_args=engine_args,

87 - usage_context=UsageContext.UNKNOWN_CONTEXT,

88 - ipc_path=ipc_path)

89 -

90 - # Raise error during first forward pass.

91 - engine.engine.model_executor.check_health = Mock(side_effect=RAISED_ERROR)

92 -

93 - # Run engine.

94 - engine.start()

95 -

96 -

97 -@pytest.mark.asyncio

98 -async def test_failed_health_check(tmp_socket):

99 - with RemoteMQLLMEngine(

100 - engine_args=ENGINE_ARGS,

101 - ipc_path=tmp_socket,

102 - run_fn=run_with_evil_model_executor_health) as engine:

103 -

104 - client = await engine.make_client()

105 - assert client.is_running

106 -

107 - # Health probe should throw RAISED_ERROR.

108 - await asyncio.sleep(15.)

109 -

110 - with pytest.raises(RAISED_ERROR):

111 - await client.check_health()

112 - assert client.errored

113 -

114 - # Generate call should throw ENGINE_DEAD_ERROR

115 - with pytest.raises(MQEngineDeadError):

116 - async for _ in client.generate(prompt="Hello my name is",

117 - sampling_params=SamplingParams(),

118 - request_id=str(uuid.uuid4())):

119 - pass

120 -

121 - client.close()

122 -

123 -

124 -def run_with_evil_abort(engine_args: AsyncEngineArgs, ipc_path: str):

125 - # Make engine.

126 - engine = MQLLMEngine.from_engine_args(

127 - engine_args=engine_args,

128 - usage_context=UsageContext.UNKNOWN_CONTEXT,

129 - ipc_path=ipc_path)

130 -

131 - # Raise error during abort call.

132 - engine.engine.abort_request = Mock(side_effect=RAISED_ERROR)

133 -

134 - # Run engine.

135 - engine.start()

136 -

137 -

138 -@pytest.mark.asyncio

139 -async def test_failed_abort(tmp_socket):

140 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

141 - ipc_path=tmp_socket,

142 - run_fn=run_with_evil_abort) as engine:

143 -

144 - client = await engine.make_client()

145 - assert client.is_running

146 -

147 - # First check health should work.

148 - await client.check_health()

149 -

150 - # Trigger an abort on the client side.

151 - # This request ID does not exist, and will cause the engine to error

152 - await client.abort(request_id="foo")

153 -

154 - # Future generation requests will now fail

155 - # with reference to the original KeyError("foo")

156 - with pytest.raises(MQEngineDeadError) as execinfo:

157 - async for _ in client.generate(

158 - prompt="Hello my name is",

159 - sampling_params=SamplingParams(max_tokens=10),

160 - request_id=str(uuid.uuid4())):

161 - pass

162 - assert "KeyError" in repr(execinfo.value)

163 - assert client.errored

164 -

165 - # This should raise the original error.

166 - with pytest.raises(RAISED_ERROR):

167 - await client.check_health()

168 -

169 - client.close()

170 -

171 -

172 -@pytest.mark.asyncio

173 -async def test_batch_error(tmp_socket):

174 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

175 - ipc_path=tmp_socket,

176 - run_fn=run_with_evil_abort) as engine:

177 -

178 - client = await engine.make_client()

179 - assert client.is_running

180 -

181 - # First check health should work.

182 - await client.check_health()

183 -

184 - # Batch of requests

185 - async def do_generate(client):

186 - # min_tokens=2048 to keep busy the engine busy

187 - # to get enough time to get process a request

188 - # that will crash the engine

189 - params = SamplingParams(min_tokens=2048, max_tokens=2048)

190 - async for _ in client.generate(prompt="Hello my name is",

191 - sampling_params=params,

192 - request_id=str(uuid.uuid4())):

193 - pass

194 -

195 - tasks = [asyncio.create_task(do_generate(client)) for _ in range(10)]

196 -

197 - # This request will force a processing batch to raise

198 - # an exception and next the engine get errored

199 - await client.abort(request_id="foo")

200 -

201 - # The batch of those request failed, then they

202 - # should get the same exception as a MQEngineDeadError.

203 - errors = await asyncio.gather(*tasks, return_exceptions=True)

204 - for e in errors:

205 - assert isinstance(e, MQEngineDeadError)

206 - assert "KeyError" in repr(e)

207 -

208 - client.close()

209 -

210 -

211 -@pytest.mark.asyncio

212 -async def test_bad_request(tmp_socket):

213 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

214 - ipc_path=tmp_socket) as engine:

215 -

216 - client = await engine.make_client()

217 -

218 - # Invalid request should fail, but not crash the server.

219 - with pytest.raises(ValueError):

220 - async for _ in client.generate(prompt="Hello my name is",

221 - sampling_params=SamplingParams(),

222 - request_id="abcd-1",

223 - lora_request=LoRARequest(

224 - "invalid-lora", 1,

225 - "invalid-path")):

226 - pass

227 -

228 - # This request should be okay.

229 - async for _ in client.generate(prompt="Hello my name is",

230 - sampling_params=SamplingParams(),

231 - request_id="abcd-2"):

232 - pass

233 -

234 - # Shutdown.

235 - client.close()

236 -

237 -

238 -@pytest.mark.asyncio

239 -async def test_mp_crash_detection(monkeypatch: pytest.MonkeyPatch):

240 - with monkeypatch.context() as m:

241 -

242 - parser = FlexibleArgumentParser(

243 - description="vLLM's remote OpenAI server.")

244 - parser = make_arg_parser(parser)

245 - args = parser.parse_args([])

246 -

247 - # When LLMEngine is loaded, it will crash.

248 - def mock_init():

249 - raise ValueError

250 -

251 - m.setattr(LLMEngine, "__init__", mock_init)

252 -

253 - start = time.perf_counter()

254 - async with build_async_engine_client(args):

255 - pass

256 - end = time.perf_counter()

257 -

258 - assert end - start < 100, (

259 - "Expected vLLM to gracefully shutdown in <100s "

260 - "if there is an error in the startup.")

261 -

262 -

263 -@pytest.mark.asyncio

264 -async def test_mp_cuda_init():

265 - # it should not crash, when cuda is initialized

266 - # in the API server process

267 - import torch

268 - torch.cuda.init()

269 - parser = FlexibleArgumentParser(description="vLLM's remote OpenAI server.")

270 - parser = make_arg_parser(parser)

271 - args = parser.parse_args([])

272 -

273 - async with build_async_engine_client(args):

274 - pass

275 -

276 -

277 -@pytest.mark.asyncio

278 -async def test_engine_process_death(tmp_socket):

279 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

280 - ipc_path=tmp_socket) as engine:

281 -

282 - client = await engine.make_client()

283 - assert client.is_running

284 -

285 - # kill the engine process

286 - engine.proc.kill()

287 -

288 - # Generate call should fail

289 - with pytest.raises(MQEngineDeadError):

290 - async for _ in client.generate(prompt="Hello my name is",

291 - sampling_params=SamplingParams(),

292 - request_id=str(uuid.uuid4())):

293 - pass

294 -

295 - # And the health check should show the engine is dead

296 - with pytest.raises(RuntimeError, match="Engine process .* died"):

297 - await client.check_health()

298 -

299 - client.close()

300 -

301 -

302 -def run_with_evil_input_processing(engine_args: AsyncEngineArgs,

303 - ipc_path: str):

304 - """Simulate an exception while preparing inputs for the model.

305 - In the wild, this could be something like a multimodal input processor

306 - failing on invalid image data."""

307 -

308 - # Make engine.

309 - engine = MQLLMEngine.from_engine_args(

310 - engine_args=engine_args,

311 - usage_context=UsageContext.UNKNOWN_CONTEXT,

312 - ipc_path=ipc_path)

313 -

314 - runner = engine.engine.model_executor.driver_worker.worker.model_runner

315 -

316 - # Raise error in the model runner when adding a sequence group.

317 - # See class ModelInputForGPUBuilder

318 - def raiser(_, seq_group_metadata: SequenceGroupMetadata):

319 - if seq_group_metadata.request_id.startswith("evil"):

320 - raise RAISED_ERROR(RAISED_VALUE)

321 -

322 - runner.builder.per_seq_group_compute_fns.append(raiser)

323 -

324 - # Run engine.

325 - engine.start()

326 -

327 -

328 -@pytest.mark.asyncio

329 -async def test_failed_inputs(tmp_socket):

330 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

331 - ipc_path=tmp_socket,

332 - run_fn=run_with_evil_input_processing) as engine:

333 -

334 - client = await engine.make_client()

335 - assert client.is_running

336 -

337 - # Engine should be healthy

338 - await client.check_health()

339 -

340 - async def run_failing_request():

341 - async for _ in client.generate(

342 - prompt="Hello my name is",

343 - sampling_params=SamplingParams(max_tokens=10),

344 - request_id="evil" + str(uuid.uuid4())):

345 - pass

346 -

347 - async def run_passing_request():

348 - async for _ in client.generate(

349 - prompt="Hello my name is",

350 - sampling_params=SamplingParams(max_tokens=10),

351 - request_id=str(uuid.uuid4())):

352 - pass

353 -

354 - passing_tasks = [

355 - asyncio.create_task(run_passing_request()) for _ in range(10)

356 - ]

357 - failing_tasks = [

358 - asyncio.create_task(run_failing_request()) for _ in range(10)

359 - ]

360 - await asyncio.gather(*failing_tasks, return_exceptions=True)

361 - await asyncio.gather(*passing_tasks)

362 -

363 - # All the bad inputs should have raised

364 - for task in failing_tasks:

365 - with pytest.raises(RAISED_ERROR):

366 - task.result()

367 -

368 - # But all good inputs should have still succeeded

369 - for task in passing_tasks:

370 - task.result()

371 -

372 - # And the engine should remain healthy

373 - assert not client.errored

374 - await client.check_health()

375 -

376 - client.close()

tests/mq_llm_engine/test_load.py deleted

-59

@@ -1,59 +0,0 @@

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

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

3 -"""Test that the MQLLMEngine is able to handle 10k concurrent requests."""

4 -

5 -import asyncio

6 -import tempfile

7 -import uuid

8 -

9 -import pytest

10 -

11 -from tests.mq_llm_engine.utils import RemoteMQLLMEngine, generate

12 -from vllm.engine.arg_utils import AsyncEngineArgs

13 -

14 -MODEL = "google/gemma-1.1-2b-it"

15 -NUM_EXPECTED_TOKENS = 10

16 -NUM_REQUESTS = 10000

17 -

18 -# Scenarios to test for num generated token.

19 -ENGINE_ARGS = AsyncEngineArgs(model=MODEL)

20 -

21 -

22 -@pytest.fixture(scope="function")

23 -def tmp_socket():

24 - with tempfile.TemporaryDirectory() as td:

25 - yield f"ipc://{td}/{uuid.uuid4()}"

26 -

27 -

28 -@pytest.mark.asyncio

29 -async def test_load(tmp_socket):

30 - with RemoteMQLLMEngine(engine_args=ENGINE_ARGS,

31 - ipc_path=tmp_socket) as engine:

32 -

33 - client = await engine.make_client()

34 -

35 - request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

36 -

37 - # Create concurrent requests.

38 - tasks = []

39 - for request_id in request_ids:

40 - tasks.append(

41 - asyncio.create_task(

42 - generate(client, request_id, NUM_EXPECTED_TOKENS)))

43 -

44 - # Confirm that we got all the EXPECTED tokens from the requests.

45 - failed_request_id = None

46 - tokens = None

47 - for task in tasks:

48 - num_generated_tokens, request_id = await task

49 - if (num_generated_tokens != NUM_EXPECTED_TOKENS

50 - and failed_request_id is None):

51 - failed_request_id = request_id

52 - tokens = num_generated_tokens

53 -

54 - assert failed_request_id is None, (

55 - f"{failed_request_id} generated {tokens} but "

56 - f"expected {NUM_EXPECTED_TOKENS}")

57 -

58 - # Shutdown.

59 - client.close()

tests/mq_llm_engine/utils.py deleted

-81

@@ -1,81 +0,0 @@

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

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

3 -

4 -import asyncio

5 -import multiprocessing

6 -from typing import Callable, Union

7 -

8 -from vllm import SamplingParams

9 -from vllm.engine.arg_utils import AsyncEngineArgs

10 -from vllm.engine.multiprocessing.client import MQLLMEngineClient

11 -from vllm.engine.multiprocessing.engine import MQLLMEngine

12 -from vllm.outputs import RequestOutput

13 -from vllm.usage.usage_lib import UsageContext

14 -

15 -

16 -async def generate(

17 - client: MQLLMEngineClient,

18 - request_id: str,

19 - num_tokens: int,

20 - return_output: bool = False) -> Union[RequestOutput, tuple[int, str]]:

21 -

22 - final_output = None

23 - count = 0

24 - async for out in client.generate(

25 - request_id=request_id,

26 - prompt="Hello my name is Robert and",

27 - sampling_params=SamplingParams(max_tokens=num_tokens,

28 - temperature=0)):

29 -

30 - count += 1

31 - final_output = out

32 - await asyncio.sleep(0.)

33 -

34 - if return_output:

35 - return final_output

36 -

37 - # Confirm we generated all the tokens we expected.

38 - return count, request_id

39 -

40 -

41 -def run_normal(engine_args: AsyncEngineArgs, ipc_path: str):

42 - # Make engine.

43 - engine = MQLLMEngine.from_engine_args(

44 - engine_args=engine_args,

45 - usage_context=UsageContext.UNKNOWN_CONTEXT,

46 - ipc_path=ipc_path)

47 -

48 - # Run engine.

49 - engine.start()

50 -

51 -

52 -class RemoteMQLLMEngine:

53 -

54 - def __init__(self,

55 - engine_args: AsyncEngineArgs,

56 - ipc_path: str,

57 - run_fn: Callable = run_normal) -> None:

58 -

59 - self.engine_args = engine_args

60 - self.ipc_path = ipc_path

61 - context = multiprocessing.get_context("spawn")

62 - self.proc = context.Process(target=run_fn,

63 - args=(engine_args, ipc_path))

64 - self.proc.start()

65 -

66 - def __enter__(self):

67 - return self

68 -

69 - def __exit__(self, exc_type, exc_value, traceback):

70 - self.proc.kill()

71 -

72 - async def make_client(self) -> MQLLMEngineClient:

73 - engine_config = self.engine_args.create_engine_config()

74 - client = MQLLMEngineClient(self.ipc_path, engine_config, self.proc.pid)

75 - while True:

76 - try:

77 - await client.setup()

78 - break

79 - except TimeoutError:

80 - assert self.proc.is_alive()

81 - return client

vllm/engine/multiprocessing/__init__.py deleted

-145

@@ -1,145 +0,0 @@

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

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

3 -

4 -import uuid

5 -from dataclasses import dataclass, field

6 -from enum import Enum

7 -from typing import List, Mapping, Optional, Union

8 -

9 -from vllm import PoolingParams

10 -from vllm.inputs import PromptType

11 -from vllm.lora.request import LoRARequest

12 -from vllm.outputs import RequestOutput

13 -from vllm.sampling_params import SamplingParams

14 -from vllm.utils import Device

15 -

16 -VLLM_RPC_SUCCESS_STR = "SUCCESS"

17 -

18 -IPC_INPUT_EXT = "_input_socket"

19 -IPC_OUTPUT_EXT = "_output_socket"

20 -IPC_HEALTH_EXT = "_health_socket"

21 -IPC_DATA_EXT = "_data_socket"

22 -

23 -

24 -class MQEngineDeadError(RuntimeError):

25 - pass

26 -

27 -

28 -@dataclass

29 -class RPCProcessRequest:

30 - prompt: PromptType

31 - params: Union[SamplingParams, PoolingParams]

32 - request_id: str

33 - lora_request: Optional[LoRARequest] = None

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

35 - priority: int = 0

36 -

37 - def __init__(

38 - self,

39 - prompt: PromptType,

40 - params: Union[SamplingParams, PoolingParams],

41 - request_id: str,

42 - lora_request: Optional[LoRARequest] = None,

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

44 - priority: int = 0,

45 - ) -> None:

46 - super().__init__()

47 -

48 - self.prompt = prompt

49 - self.params = params

50 - self.request_id = request_id

51 - self.lora_request = lora_request

52 - self.trace_headers = trace_headers

53 - self.priority = priority

54 -

55 -

56 -@dataclass

57 -class RPCError:

58 - request_id: Optional[str]

59 - is_engine_errored: bool

60 - exception: BaseException

61 -

62 -

63 -@dataclass

64 -class RPCAbortRequest:

65 - request_id: str

66 -

67 -

68 -class RPCStartupRequest(Enum):

69 - IS_SERVER_READY = 1

70 -

71 -

72 -@dataclass

73 -class RPCStartupResponse:

74 - tracing_enabled: bool

75 -

76 -

77 -class RPCUProfileRequest(Enum):

78 - START_PROFILE = 1

79 - STOP_PROFILE = 2

80 -

81 -

82 -class RPCResetMultiModalCacheRequest(Enum):

83 - RESET = 1

84 -

85 -

86 -@dataclass

87 -class RPCResetPrefixCacheRequest:

88 - device: Device

89 -

90 -

91 -class RPCSleepRequest(Enum):

92 - SLEEP_LEVEL_1 = 1

93 - SLEEP_LEVEL_2 = 2

94 -

95 -

96 -@dataclass

97 -class RPCWakeUpRequest:

98 - tags: Optional[list[str]] = None

99 -

100 -

101 -@dataclass

102 -class RPCIsSleepingRequest:

103 - # Set the default value of request_id to a new UUID

104 - request_id: str = field(default_factory=lambda: str(uuid.uuid4()))

105 -

106 -

107 -@dataclass

108 -class RPCIsSleepingResponse:

109 - request_id: str

110 - is_sleeping: bool

111 -

112 -

113 -@dataclass

114 -class RPCLoadAdapterRequest:

115 - lora_request: LoRARequest

116 - # Set the default value of request_id to a new UUID

117 - request_id: str = field(default_factory=lambda: str(uuid.uuid4()))

118 -

119 -

120 -@dataclass

121 -class RPCAdapterLoadedResponse:

122 - request_id: str

123 - lora_loaded: bool

124 -

125 -

126 -RPC_REQUEST_T = Union[RPCProcessRequest, RPCAbortRequest, RPCStartupRequest,

127 - RPCUProfileRequest, RPCLoadAdapterRequest,

128 - RPCResetMultiModalCacheRequest,

129 - RPCResetPrefixCacheRequest, RPCSleepRequest,

130 - RPCWakeUpRequest, RPCIsSleepingRequest]

131 -

132 -REQUEST_OUTPUTS_T = Union[List[RequestOutput], RPCAdapterLoadedResponse,

133 - RPCIsSleepingResponse, RPCError]

134 -

135 -

136 -def ENGINE_DEAD_ERROR(

137 - error: Optional[BaseException] = None) -> MQEngineDeadError:

138 - if error is None:

139 - return MQEngineDeadError(

140 - "Engine loop is not running. Inspect the stacktrace to "

141 - "find the original error")

142 -

143 - return MQEngineDeadError(

144 - "Engine loop is not running. Inspect the stacktrace to "

145 - f"find the original error: {repr(error)}.")

vllm/engine/multiprocessing/client.py deleted

-643

@@ -1,643 +0,0 @@

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

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

3 -

4 -import asyncio

5 -import copy

6 -import pickle

7 -from contextlib import contextmanager, suppress

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

9 - Mapping, Optional, Union)

10 -

11 -import cloudpickle

12 -import psutil

13 -import zmq

14 -import zmq.asyncio

15 -from zmq import Frame # type: ignore[attr-defined]

16 -from zmq.asyncio import Socket

17 -

18 -from vllm import PoolingParams

19 -from vllm.config import DecodingConfig, ModelConfig, VllmConfig

20 -from vllm.core.scheduler import SchedulerOutputs

21 -# yapf conflicts with isort for this block

22 -# yapf: disable

23 -from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

24 - IPC_HEALTH_EXT, IPC_INPUT_EXT,

25 - IPC_OUTPUT_EXT, RPC_REQUEST_T,

26 - VLLM_RPC_SUCCESS_STR, RPCAbortRequest,

27 - RPCAdapterLoadedResponse, RPCError,

28 - RPCIsSleepingRequest,

29 - RPCIsSleepingResponse,

30 - RPCLoadAdapterRequest,

31 - RPCProcessRequest,

32 - RPCResetMultiModalCacheRequest,

33 - RPCResetPrefixCacheRequest,

34 - RPCSleepRequest, RPCStartupRequest,

35 - RPCStartupResponse,

36 - RPCUProfileRequest, RPCWakeUpRequest)

37 -from vllm.engine.protocol import EngineClient

38 -# yapf: enable

39 -from vllm.envs import VLLM_RPC_TIMEOUT

40 -from vllm.inputs import PromptType

41 -from vllm.inputs.preprocess import InputPreprocessor

42 -from vllm.logger import init_logger

43 -from vllm.lora.request import LoRARequest

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

45 -from vllm.outputs import PoolingRequestOutput, RequestOutput

46 -from vllm.sampling_params import SamplingParams

47 -from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs

48 -from vllm.utils import Device

49 -

50 -logger = init_logger(__name__)

51 -

52 -

53 -class MQClientClosedError(Exception):

54 - """Exception class raised when the client is used post-close.

55 -

56 - The client can be closed, which closes the ZMQ context. This normally

57 - happens on server shutdown. In some cases, methods like abort and

58 - do_log_stats will still be called and then try to open a socket, which

59 - causes a ZMQError and creates a huge stack trace.

60 - So, we throw this error such that we can suppress it.

61 - """

62 -

63 -

64 -class MQLLMEngineClient(EngineClient):

65 - """A client wrapper for MQLLMEngine that conforms to the

66 - EngineClient protocol.

67 -

68 - MQLLMEngine and MQLLMEngineClient are intended to run in separate

69 - processes communicating via zeromq ipc sockets.

70 -

71 - The entrypoint to MQLLMEngineClient is through the generate()

72 - method. On generate() MQLLMEngine does three things:

73 - - Creates an asyncio output queue

74 - - Sends a RPCGenerateRequest to the MQLLMEngine via zmq

75 - - Pulls RequestOutputs from its queue and yields them

76 -

77 - MQLLMEngine runs two background loops:

78 - - output_loop: the output loop pulls List[RequestOutput]

79 - from the MQLLMEngine via zmq (each list is the output

80 - of one engine_step in the LLMEngine). It then parses

81 - the list and pushes individual request_outputs into

82 - the corresponding output_queue such that they can be

83 - consumed by the .generate() method.

84 - - health_loop: the health loop queries the health socket

85 - every N seconds, confirming the engine is healthy

86 - """

87 -

88 - def __init__(self, ipc_path: str, engine_config: VllmConfig,

89 - engine_pid: int):

90 - self.context = zmq.asyncio.Context()

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

92 -

93 - # Get the configs.

94 - self.vllm_config = engine_config

95 - self.model_config = engine_config.model_config

96 - self.decoding_config = engine_config.decoding_config

97 -

98 - if self.vllm_config.model_config.skip_tokenizer_init:

99 - self.tokenizer = None

100 -

101 - else:

102 - # Create the tokenizer group.

103 - self.tokenizer = init_tokenizer_from_configs(

104 - model_config=self.model_config,

105 - scheduler_config=engine_config.scheduler_config,

106 - lora_config=engine_config.lora_config)

107 -

108 - self.input_preprocessor = InputPreprocessor(self.model_config,

109 - self.tokenizer)

110 -

111 - # Send RPCGenerateRequest to the MQLLMEngine.

112 - self.input_socket: Socket = self.context.socket(zmq.constants.PUSH)

113 - self.input_socket.connect(f"{ipc_path}{IPC_INPUT_EXT}")

114 -

115 - # Receive streams of RequestOutput from the MQLLMEngine.

116 - self.output_socket: Socket = self.context.socket(zmq.constants.PULL)

117 - self.output_socket.connect(f"{ipc_path}{IPC_OUTPUT_EXT}")

118 -

119 - # IPC path for acking heartbeats.

120 - self.heartbeat_socket: Socket = self.context.socket(zmq.constants.PULL)

121 - self.heartbeat_socket.connect(f"{ipc_path}{IPC_HEALTH_EXT}")

122 -

123 - # IPC path for the data socket.

124 - self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"

125 -

126 - # Stream for each individual request.

127 - self.output_queues: Dict[str, asyncio.Queue] = {}

128 -

129 - # Loop to handle output of the LLMEngine periodically.

130 - # Started after the MQLLMEngine is ready so that we can

131 - # build the Client in an executor to enable clean shutdown.

132 - self.output_loop: Optional[asyncio.Task] = None

133 -

134 - # Loop to check health of the LLMEngine periodically.

135 - # Started after the MQLLMEngine is ready.

136 - self.health_loop: Optional[asyncio.Task] = None

137 - self._engine_process = psutil.Process(engine_pid)

138 -

139 - @staticmethod

140 - def is_unsupported_config(vllm_config: VllmConfig):

141 - # Pipeline parallel not yet supported

142 - return vllm_config.parallel_config.pipeline_parallel_size > 1

143 -

144 - @contextmanager

145 - def get_data_socket(self) -> Iterator[Socket]:

146 - socket = self.context.socket(zmq.constants.DEALER)

147 - try:

148 - socket.connect(self.data_ipc_path)

149 - yield socket

150 - finally:

151 - socket.close(linger=0)

152 -

153 - async def run_heartbeat_loop(self, timeout: int):

154 - """Background loop that continually checks to ensure the engine process

155 - is still alive.

156 - """

157 - try:

158 - while True:

159 - # Check if the engine process is running:

160 - if not self._engine_process.is_running() or (

161 - self._engine_process.status() == psutil.STATUS_ZOMBIE):

162 - # NB: is_running() returns True for zombies

163 - self._set_errored(

164 - RuntimeError(

165 - f"Engine process (pid {self._engine_process.pid}) "

166 - "died."))

167 - break

168 -

169 - if await self.heartbeat_socket.poll(timeout=timeout):

170 - # Heartbeat received- check the message

171 - await self._check_success(

172 - error_message="Heartbeat failed.",

173 - socket=self.heartbeat_socket)

174 -

175 - logger.debug("Heartbeat successful.")

176 -

177 - except asyncio.CancelledError:

178 - logger.debug("Shutting down MQLLMEngineClient check health loop.")

179 -

180 - except psutil.NoSuchProcess:

181 - self._set_errored(

182 - RuntimeError(

183 - f"Engine process (pid {self._engine_process.pid}) died."))

184 -

185 - except Exception as e:

186 - self._set_errored(e)

187 -

188 - async def run_output_handler_loop(self):

189 - """Get RequestOutputs from Engine and stream to Request Queues"""

190 -

191 - try:

192 - while True:

193 - # Poll, checking for ENGINE_DEAD

194 - while await self.output_socket.poll(timeout=VLLM_RPC_TIMEOUT

195 - ) == 0:

196 - logger.debug("Waiting for output from MQLLMEngine.")

197 -

198 - # If errored, alert all running requests.

199 - if self.errored:

200 - for queue_j in tuple(self.output_queues.values()):

201 - queue_j.put_nowait(

202 - ENGINE_DEAD_ERROR(self._errored_with))

203 - return

204 -

205 - message: Frame = await self.output_socket.recv(copy=False)

206 - request_outputs = pickle.loads(message.buffer)

207 -

208 - is_error = isinstance(request_outputs,

209 - (BaseException, RPCError))

210 - if is_error:

211 - if isinstance(request_outputs, RPCError):

212 - rpc_error: RPCError = request_outputs

213 - request_id = rpc_error.request_id

214 - exception = rpc_error.exception

215 - is_engine_errored = rpc_error.is_engine_errored

216 - else:

217 - # MPLLMEngine should always return an RPCError to

218 - # the output_socket when an issue arises.

219 - # If we are here, we are in a bad state and

220 - # should shut down the server.

221 - error: BaseException = request_outputs

222 - logger.error(

223 - "Received Exception %s rather than RPCError from "

224 - "MPLLMEngine. This should never happen.", error)

225 - request_id = None

226 - exception = error

227 - is_engine_errored = True

228 -

229 - # Set to error state only on engine critical error

230 - # (and record only the first one)

231 - if is_engine_errored and not self._errored_with:

232 - self._errored_with = exception

233 - # If engine is errored, no matter the type of exception

234 - # it will no longer be able to receive new requests,

235 - # therefore we have to inform that the current

236 - # processed requests failed as well. Send back a dead

237 - # engine error give this feedback and also give a

238 - # 'hint' to the server to shut down next.

239 - exception = self.dead_error

240 -

241 - if request_id is None:

242 - # If request_id is None, then the engine raised an

243 - # exception for a batch, and we may not know the

244 - # request that caused it, neither if it was actually

245 - # caused by any of them (e.g. CUDA OOM). Therefore we

246 - # broadcast the same exception for all requests.

247 - for queue_i in tuple(self.output_queues.values()):

248 - queue_i.put_nowait(exception)

249 - else:

250 - queue = self.output_queues.get(request_id)

251 - if queue is not None:

252 - queue.put_nowait(exception)

253 - # Put each output into the appropriate queue.

254 - elif isinstance(

255 - request_outputs,

256 - (RPCAdapterLoadedResponse, RPCIsSleepingResponse)):

257 - self._add_output(request_outputs)

258 - else:

259 - for request_output in request_outputs:

260 - self._add_output(request_output)

261 -

262 - except asyncio.CancelledError:

263 - logger.debug("Shutting down MQLLMEngineClient output handler.")

264 -

265 - def _add_output(self, request_output: Union[RequestOutput,

266 - RPCAdapterLoadedResponse,

267 - RPCIsSleepingResponse]):

268 - queue = self.output_queues.get(request_output.request_id)

269 - if queue is not None:

270 - queue.put_nowait(request_output)

271 -

272 - async def setup(self):

273 - """Set up the client before it starts sending server requests."""

274 -

275 - # Start output_loop

276 - if self.output_loop is None:

277 - # only generate once to avoid multiple concurrent output_loops

278 - # this will lead to race conditions and wrong orders of tokens

279 - # returned by the engine

280 - # setup will be called multiple times during the startup of

281 - # the engine

282 - self.output_loop = asyncio.create_task(

283 - self.run_output_handler_loop())

284 -

285 - with self.get_data_socket() as socket:

286 - # Wait until server is ready.

287 - response = await self._wait_for_server_rpc(socket)

288 -

289 - self.tracing_flag = response.tracing_enabled

290 -

291 - # Start health_loop.

292 - if self.health_loop is None:

293 - self.health_loop = asyncio.create_task(

294 - self.run_heartbeat_loop(timeout=VLLM_RPC_TIMEOUT))

295 -

296 - def close(self):

297 - """Destroy the ZeroMQ Context."""

298 - # Close all sockets and terminate the context.

299 - self.context.destroy(linger=0)

300 -

301 - # Cancel background tasks.

302 - if self.health_loop is not None:

303 - self.health_loop.cancel()

304 - if self.output_loop is not None:

305 - self.output_loop.cancel()

306 -

307 - def _set_errored(self, e: BaseException):

308 - logger.exception(repr(e))

309 - if self._errored_with is None:

310 - self._errored_with = e

311 -

312 - @staticmethod

313 - async def _send_get_data_rpc_request(request: RPCStartupRequest,

314 - expected_type: Any,

315 - error_message: str,

316 - socket: Socket) -> Any:

317 - """Send an RPC request that is expecting data back."""

318 -

319 - # Ping RPCServer with a request.

320 - await socket.send_multipart((pickle.dumps(request), ), copy=False)

321 -

322 - # Make sure the server responds in time.

323 - if await socket.poll(timeout=VLLM_RPC_TIMEOUT) == 0:

324 - raise TimeoutError("RPCServer didn't reply within "

325 - f"{VLLM_RPC_TIMEOUT} ms")

326 -

327 - # Await the data from the Server.

328 - frame = await socket.recv(copy=False)

329 - data = pickle.loads(frame.buffer)

330 -

331 - if isinstance(data, BaseException):

332 - raise data

333 - elif not isinstance(data, expected_type):

334 - raise ValueError(error_message)

335 -

336 - return data

337 -

338 - @staticmethod

339 - async def _send_one_way_rpc_request(request: RPC_REQUEST_T,

340 - socket: Socket):

341 - """Send one-way RPC request to trigger an action."""

342 -

343 - if socket.closed:

344 - raise MQClientClosedError()

345 -

346 - await socket.send_multipart((pickle.dumps(request), ))

347 -

348 - async def _await_ack(self, error_message: str, socket: Socket):

349 - """Await acknowledgement that a request succeeded."""

350 -

351 - if socket.closed:

352 - raise MQClientClosedError()

353 -

354 - if await socket.poll(timeout=VLLM_RPC_TIMEOUT) == 0:

355 - raise TimeoutError("MQLLMEngine didn't reply within "

356 - f"{VLLM_RPC_TIMEOUT}ms")

357 -

358 - await self._check_success(error_message, socket)

359 -

360 - @staticmethod

361 - async def _check_success(error_message: str, socket: Socket):

362 - """Confirm that socket has a VLLM_RPC_SUCCESS_STR message"""

363 -

364 - if socket.closed:

365 - raise MQClientClosedError()

366 -

367 - frame = await socket.recv(copy=False)

368 - response = pickle.loads(frame.buffer)

369 -

370 - # Raise error if unsuccessful

371 - if isinstance(response, BaseException):

372 - raise response

373 - elif (not isinstance(response, str)

374 - or response != VLLM_RPC_SUCCESS_STR):

375 - raise ValueError(error_message)

376 -

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

378 - return self.input_preprocessor

379 -

380 - async def get_tokenizer(self, lora_request: Optional[LoRARequest] = None):

381 - if self.tokenizer is None:

382 - return None

383 - else:

384 - return await self.tokenizer.get_lora_tokenizer_async(lora_request)

385 -

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

387 - return self.vllm_config

388 -

389 - async def get_decoding_config(self) -> DecodingConfig:

390 - return self.decoding_config

391 -

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

393 - return self.model_config

394 -

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

396 - return self.tracing_flag

397 -

398 - async def _wait_for_server_rpc(self, socket: Socket) -> RPCStartupResponse:

399 - """Wait for the RPCServer to start up."""

400 -

401 - return await self._send_get_data_rpc_request(

402 - request=RPCStartupRequest.IS_SERVER_READY,

403 - expected_type=RPCStartupResponse,

404 - error_message="Unable to start RPC Server",

405 - socket=socket)

406 -

407 - async def abort(self, request_id: Union[str, Iterable[str]]):

408 - """Send an ABORT_REQUEST signal to the RPC Server"""

409 -

410 - if not isinstance(request_id, str):

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

412 - " deprecated V0")

413 -

414 - with suppress(MQClientClosedError):

415 - await self._send_one_way_rpc_request(

416 - request=RPCAbortRequest(request_id), socket=self.input_socket)

417 -

418 - async def do_log_stats(

419 - self,

420 - scheduler_outputs: Optional[SchedulerOutputs] = None,

421 - model_output: Optional[List[SamplerOutput]] = None,

422 - ) -> None:

423 - """

424 - Ignore do_log_stats (handled on MQLLMEngine polling)

425 - """

426 - pass

427 -

428 - async def check_health(self):

429 - """

430 - The check health loop probes the health status of the

431 - Engine's health every N seconds and sets _errored_with

432 - if the engine is unhealthy.

433 - """

434 - if self._errored_with is not None:

435 - raise self._errored_with

436 -

437 - @property

438 - def is_running(self) -> bool:

439 - return not self.errored

440 -

441 - @property

442 - def is_stopped(self) -> bool:

443 - return self.errored

444 -

445 - @property

446 - def errored(self) -> bool:

447 - return self._errored_with is not None

448 -

449 - @property

450 - def dead_error(self) -> BaseException:

451 - return ENGINE_DEAD_ERROR(self._errored_with)

452 -

453 - def generate(

454 - self,

455 - prompt: PromptType,

456 - sampling_params: SamplingParams,

457 - request_id: str,

458 - lora_request: Optional[LoRARequest] = None,

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

460 - priority: int = 0,

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

462 - """Generate outputs for a request.

463 -

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

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

466 - from the LLMEngine to the caller.

467 -

468 - Args:

469 - prompt: The prompt to the LLM. See

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

471 - the format of each input.

472 - sampling_params: The sampling parameters of the request.

473 - request_id: The unique id of the request.

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

475 - trace_headers: OpenTelemetry trace headers.

476 - priority: Priority of the request (lower means earlier handling).

477 - Any priority other than 0 will lead to an error if the

478 - scheduling policy is not "priority".

479 - """

480 - return self._process_request(prompt, sampling_params, request_id,

481 - lora_request, trace_headers, priority)

482 -

483 - def encode(

484 - self,

485 - prompt: PromptType,

486 - pooling_params: PoolingParams,

487 - request_id: str,

488 - lora_request: Optional[LoRARequest] = None,

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

490 - priority: int = 0,

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

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

493 - raise NotImplementedError(

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

495 -

496 - async def _process_request(

497 - self,

498 - prompt: PromptType,

499 - params: SamplingParams,

500 - request_id: str,

501 - lora_request: Optional[LoRARequest] = None,

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

503 - priority: int = 0,

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

505 - """Send an RPCGenerateRequest to the RPCServer and stream responses."""

506 -

507 - # If already dead, error out.

508 - if self._errored_with is not None:

509 - raise ENGINE_DEAD_ERROR(self._errored_with)

510 -

511 - # Ensure the request id is unique among running requests

512 - if request_id in self.output_queues:

513 - raise ValueError(f"Request {request_id} already exists")

514 -

515 - # 1) Create output queue for this request.

516 - queue: asyncio.Queue[Union[RequestOutput,

517 - BaseException]] = asyncio.Queue()

518 - self.output_queues[request_id] = queue

519 -

520 - try:

521 - # 2) Detach logits processors so that they can be pickled

522 - # separately (may require cloudpickle which is slower)

523 - if params.logits_processors:

524 - # Defensive shallow copy

525 - params = copy.copy(params)

526 - logits_processors = params.logits_processors

527 - params.logits_processors = None

528 - lp_bytes = cloudpickle.dumps(logits_processors)

529 - else:

530 - lp_bytes = None

531 -

532 - request_bytes = pickle.dumps(

533 - RPCProcessRequest(

534 - prompt=prompt,

535 - params=params,

536 - request_id=request_id,

537 - lora_request=lora_request,

538 - trace_headers=trace_headers,

539 - priority=priority,

540 - ))

541 -

542 - # 3) Send the RPCGenerateRequest to the MQLLMEngine.

543 - parts = (request_bytes,

544 - lp_bytes) if lp_bytes else (request_bytes, )

545 - await self.input_socket.send_multipart(parts, copy=False)

546 -

547 - # 4) Stream the RequestOutputs from the output queue. Note

548 - # that the output_loop pushes RequestOutput objects to this

549 - # queue after pulling them from the zmq socket.

550 - finished = False

551 - try:

552 - while not finished:

553 - request_output = await queue.get()

554 -

555 - if isinstance(request_output, BaseException):

556 - raise request_output

557 -

558 - finished = request_output.finished

559 - yield request_output

560 - finally:

561 - # Request was canceled by the client.

562 - if not finished and not self.errored:

563 - await self.abort(request_id)

564 - finally:

565 - self.output_queues.pop(request_id)

566 -

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

568 - """Start profiling the engine"""

569 -

570 - await self._send_one_way_rpc_request(

571 - request=RPCUProfileRequest.START_PROFILE, socket=self.input_socket)

572 -

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

574 - """Stop profiling the engine"""

575 -

576 - await self._send_one_way_rpc_request(

577 - request=RPCUProfileRequest.STOP_PROFILE, socket=self.input_socket)

578 -

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

580 - """Reset the multi-modal cache"""

581 -

582 - await self._send_one_way_rpc_request(

583 - request=RPCResetMultiModalCacheRequest.RESET,

584 - socket=self.input_socket)

585 -

586 - async def reset_prefix_cache(self,

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

588 - """Reset the prefix cache"""

589 -

590 - await self._send_one_way_rpc_request(

591 - request=RPCResetPrefixCacheRequest(device),

592 - socket=self.input_socket)

593 -

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

595 - """Sleep the engine for a given level"""

596 - return await self._send_one_way_rpc_request(

597 - request=RPCSleepRequest(level), socket=self.input_socket)

598 -

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

600 - """Wake up the engine"""

601 - return await self._send_one_way_rpc_request(

602 - request=RPCWakeUpRequest(tags), socket=self.input_socket)

603 -

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

605 - """Check whether the engine is sleeping"""

606 - request = RPCIsSleepingRequest()

607 -

608 - queue: asyncio.Queue[Union[BaseException,

609 - RPCIsSleepingResponse]] = asyncio.Queue()

610 - self.output_queues[request.request_id] = queue

611 -

612 - request_bytes = pickle.dumps(request)

613 - await self.input_socket.send_multipart((request_bytes, ), copy=False)

614 -

615 - request_output = await queue.get()

616 - self.output_queues.pop(request.request_id)

617 -

618 - if isinstance(request_output, BaseException):

619 - raise request_output

620 - return request_output.is_sleeping

621 -

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

623 - """Load a new LoRA adapter into the engine for future requests."""

624 - # Uses the same I/O as generate requests

625 - request = RPCLoadAdapterRequest(lora_request)

626 -

627 - # Create output queue for this request.

628 - queue: asyncio.Queue[Union[

629 - BaseException, RPCAdapterLoadedResponse]] = asyncio.Queue()

630 - self.output_queues[request.request_id] = queue

631 -

632 - # Send the request

633 - request_bytes = pickle.dumps(request)

634 - await self.input_socket.send_multipart((request_bytes, ), copy=False)

635 -

636 - # Wait for the response

637 - request_output = await queue.get()

638 - self.output_queues.pop(request.request_id)

639 -

640 - # Raise on error, otherwise happily return None

641 - if isinstance(request_output, BaseException):

642 - raise request_output

643 - return request_output.lora_loaded

vllm/engine/multiprocessing/engine.py deleted

-470

@@ -1,470 +0,0 @@

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

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

3 -

4 -import pickle

5 -import signal

6 -from contextlib import contextmanager

7 -from typing import Iterator, List, Optional, Union

8 -

9 -import cloudpickle

10 -import zmq

11 -

12 -from vllm import AsyncEngineArgs, SamplingParams

13 -from vllm.config import VllmConfig

14 -from vllm.engine.llm_engine import LLMEngine

15 -# yapf conflicts with isort for this block

16 -# yapf: disable

17 -from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

18 - IPC_HEALTH_EXT, IPC_INPUT_EXT,

19 - IPC_OUTPUT_EXT, REQUEST_OUTPUTS_T,

20 - VLLM_RPC_SUCCESS_STR, RPCAbortRequest,

21 - RPCAdapterLoadedResponse, RPCError,

22 - RPCIsSleepingRequest,

23 - RPCIsSleepingResponse,

24 - RPCLoadAdapterRequest,

25 - RPCProcessRequest,

26 - RPCResetMultiModalCacheRequest,

27 - RPCResetPrefixCacheRequest,

28 - RPCSleepRequest, RPCStartupRequest,

29 - RPCStartupResponse,

30 - RPCUProfileRequest, RPCWakeUpRequest)

31 -# yapf: enable

32 -from vllm.logger import init_logger

33 -from vllm.outputs import RequestOutput

34 -from vllm.transformers_utils.config import (

35 - maybe_register_config_serialize_by_value)

36 -from vllm.usage.usage_lib import UsageContext

37 -from vllm.utils import deprecate_kwargs

38 -from vllm.worker.model_runner_base import InputProcessingError

39 -

40 -logger = init_logger(__name__)

41 -

42 -POLLING_TIMEOUT_MS = 10000

43 -HEALTHY_RESPONSE = (pickle.dumps(VLLM_RPC_SUCCESS_STR), )

44 -

45 -

46 -class MQLLMEngine:

47 - """A multiprocessing wrapper for

48 - [`LLMEngine`][vllm.engine.llm_engine.LLMEngine].

49 -

50 - This class is used to wrap the

51 - [`LLMEngine`][vllm.engine.llm_engine.LLMEngine] class to enable use

52 - in concurrent manner. It runs a background loop and uses zeromq to

53 - receive new requests and stream outputs incrementally via ipc.

54 -

55 - The [`LLMEngine`][vllm.engine.llm_engine.LLMEngine] generate or encode

56 - process is kicked off when a new RPCProcessRequest is received by the

57 - input_socket.

58 -

59 - The self.engine_loop checks the input_socket for new requests,

60 - adds them to the LLMEngine if there are any, calls the internal

61 - [`LLMEngine.step()`][vllm.engine.llm_engine.LLMEngine.step], and sends

62 - the RequestOutputs back over the output_socket.

63 -

64 - If use_async_sockets is set, the logic associated with reading new

65 - requests from the socket and sending data to the socket is passed

66 - as a callback to the llm_engine, which calls the logic asynchronously

67 - such that the IPC can be overlapped with the GPU.

68 -

69 - Args:

70 - ipc_path: Base path for zeromq interprocess messaging

71 - use_async_sockets: Whether to make send/recv async with GPU

72 - log_requests: Whether to log the requests.

73 - *args: Arguments for [`LLMEngine`][vllm.engine.llm_engine.LLMEngine].

74 - **kwargs: Arguments for [`LLMEngine`][vllm.engine.llm_engine.LLMEngine].

75 - """

76 -

77 - def __init__(self,

78 - ipc_path: str,

79 - use_async_sockets: bool,

80 - *args,

81 - log_requests: bool = True,

82 - **kwargs) -> None:

83 - # For MQLLMEngine, we can use cached outputs, since each new request

84 - # output is immediately pickled and send over the socket, which frees

85 - # the python object to be reused again.

86 - kwargs['use_cached_outputs'] = True

87 -

88 - self.engine = LLMEngine(*args, **kwargs)

89 - self.log_requests = log_requests

90 -

91 - self.use_async_sockets = use_async_sockets

92 - if self.use_async_sockets:

93 - self.engine.process_request_outputs_callback = \

94 - self._async_socket_engine_callback

95 -

96 - self.ctx = zmq.Context() # type: ignore[attr-defined]

97 -

98 - # Receive input from the client.

99 - self.input_socket = self.ctx.socket(zmq.constants.PULL)

100 - self.input_socket.bind(f"{ipc_path}{IPC_INPUT_EXT}")

101 -

102 - # Send output stream back to client.

103 - self.output_socket = self.ctx.socket(zmq.constants.PUSH)

104 - self.output_socket.bind(f"{ipc_path}{IPC_OUTPUT_EXT}")

105 -

106 - # Send heartbeats back to client.

107 - self.heartbeat_socket = self.ctx.socket(zmq.constants.PUSH)

108 - self.heartbeat_socket.bind(f"{ipc_path}{IPC_HEALTH_EXT}")

109 -

110 - # IPC path for the data socket.

111 - self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"

112 -

113 - # Error state.

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

115 -

116 - @property

117 - def dead_error(self) -> BaseException:

118 - if self._errored_with is not None:

119 - return ENGINE_DEAD_ERROR(self._errored_with)

120 - else:

121 - return ENGINE_DEAD_ERROR()

122 -

123 - @classmethod

124 - @deprecate_kwargs(

125 - "disable_log_requests",

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

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

128 - )

129 - def from_vllm_config(

130 - cls,

131 - vllm_config: VllmConfig,

132 - usage_context: UsageContext,

133 - enable_log_requests: bool,

134 - disable_log_stats: bool,

135 - ipc_path: str,

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

137 - ) -> "MQLLMEngine":

138 - # Setup plugins for each process

139 - from vllm.plugins import load_general_plugins

140 - load_general_plugins()

141 -

142 - use_async_sockets = vllm_config.model_config.use_async_output_proc

143 -

144 - return cls(

145 - vllm_config=vllm_config,

146 - executor_class=LLMEngine._get_executor_cls(vllm_config),

147 - ipc_path=ipc_path,

148 - usage_context=usage_context,

149 - use_async_sockets=use_async_sockets,

150 - log_requests=enable_log_requests,

151 - log_stats=(not disable_log_stats),

152 - )

153 -

154 - @staticmethod

155 - def from_engine_args(engine_args: AsyncEngineArgs,

156 - usage_context: UsageContext, ipc_path: str):

157 - """Creates an MQLLMEngine from the engine arguments."""

158 -

159 - vllm_config = engine_args.create_engine_config(usage_context)

160 - return MQLLMEngine.from_vllm_config(

161 - ipc_path=ipc_path,

162 - vllm_config=vllm_config,

163 - usage_context=usage_context,

164 - enable_log_requests=engine_args.enable_log_requests,

165 - disable_log_stats=engine_args.disable_log_stats,

166 - )

167 -

168 - def start(self):

169 - try:

170 - try:

171 - logger.debug("Starting Startup Loop.")

172 - self.run_startup_loop()

173 - logger.debug("Starting Engine Loop.")

174 - self.run_engine_loop()

175 - except Exception as e:

176 - logger.exception(repr(e))

177 - except KeyboardInterrupt:

178 - logger.debug("Shutting down MQLLMEngine.")

179 - finally:

180 - logger.debug("MQLLMEngine is shut down.")

181 - self.cleanup()

182 -

183 - def cleanup(self):

184 - """Cleanup zeromq state on shutdown."""

185 - # Closes all sockets and destroys context.

186 - self.ctx.destroy(linger=0)

187 - del self.engine

188 -

189 - @contextmanager

190 - def make_data_socket(

191 - self) -> Iterator[zmq.Socket]: # type: ignore[name-defined]

192 - socket = self.ctx.socket(zmq.constants.ROUTER)

193 - try:

194 - socket.bind(self.data_ipc_path)

195 - yield socket

196 - finally:

197 - socket.close(linger=0)

198 -

199 - def run_startup_loop(self) -> None:

200 - """Startup loop for sending data from Engine -> Client."""

201 -

202 - with self.make_data_socket() as socket:

203 - response: Union[RPCStartupResponse, BaseException]

204 - try:

205 - identity, message = socket.recv_multipart(copy=False)

206 - request: RPCStartupRequest = pickle.loads(message.buffer)

207 -

208 - # Handle the query from the Client.

209 - if request == RPCStartupRequest.IS_SERVER_READY:

210 - tracing_enabled = self.engine.is_tracing_enabled()

211 - response = RPCStartupResponse(

212 - tracing_enabled=tracing_enabled)

213 -

214 - except Exception as e:

215 - response = e

216 -

217 - socket.send_multipart((identity, pickle.dumps(response)),

218 - copy=False)

219 -

220 - def run_engine_loop(self):

221 - """Core busy loop of the LLMEngine."""

222 -

223 - while True:

224 - if not self.engine.has_unfinished_requests():

225 - # Poll until there is work to do.

226 - while self.input_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:

227 - # When there's no work, check on engine health and send

228 - # health status back to client

229 - self._health_check()

230 - self.engine.do_log_stats()

231 - logger.debug("Waiting for new requests in engine loop.")

232 -

233 - # Handle any input from the client.

234 - self.handle_new_input()

235 -

236 - # Engine step.

237 - request_outputs = self.engine_step()

238 -

239 - # Send request outputs (if async, done in engine_step callback).

240 - if not self.use_async_sockets:

241 - self._send_outputs(request_outputs)

242 -

243 - def engine_step(self) -> List[RequestOutput]:

244 - """Engine step wrapper with error handling."""

245 - try:

246 - return self.engine.step()

247 - except SystemExit:

248 - raise

249 - except InputProcessingError as e:

250 - # Special case where we handle an error preparing the inputs for

251 - # a single request in the batch

252 - rpc_err = RPCError(request_id=e.request_id,

253 - is_engine_errored=False,

254 - exception=e.__cause__)

255 - self._send_outputs(rpc_err)

256 - return []

257 - except BaseException as e:

258 - self._set_errored(e)

259 - rpc_err = RPCError(request_id=None,

260 - is_engine_errored=True,

261 - exception=e)

262 - self._send_outputs(rpc_err)

263 - raise e

264 -

265 - def handle_new_input(self):

266 - """Handle new input from the socket"""

267 - try:

268 - while self.input_socket.poll(timeout=0) != 0:

269 - frames = self.input_socket.recv_multipart(copy=False)

270 - request = pickle.loads(frames[0].buffer)

271 -

272 - if isinstance(request, RPCProcessRequest):

273 - if len(frames) > 1:

274 - # Use cloudpickle for logits processors

275 - assert isinstance(request.params, SamplingParams)

276 - lprocs = cloudpickle.loads(frames[1].buffer)

277 - request.params.logits_processors = lprocs

278 - self._handle_process_request(request)

279 - elif isinstance(request, RPCAbortRequest):

280 - self._handle_abort_request(request)

281 - elif isinstance(request, RPCUProfileRequest):

282 - if request == RPCUProfileRequest.START_PROFILE:

283 - self.start_profile()

284 - else:

285 - self.stop_profile()

286 - elif isinstance(request, RPCLoadAdapterRequest):

287 - self._handle_load_adapter_request(request)

288 - elif isinstance(request, RPCResetMultiModalCacheRequest):

289 - self.reset_mm_cache()

290 - elif isinstance(request, RPCResetPrefixCacheRequest):

291 - self.reset_prefix_cache()

292 - elif isinstance(request, RPCSleepRequest):

293 - self.sleep(request.value)

294 - elif isinstance(request, RPCWakeUpRequest):

295 - self.wake_up(request.tags)

296 - elif isinstance(request, RPCIsSleepingRequest):

297 - self._handle_is_sleeping_request(request)

298 - else:

299 - raise ValueError("Unknown RPCRequest Type: "

300 - f"{type(request)}")

301 -

302 - except Exception as e:

303 - self._set_errored(e)

304 - self._send_unhealthy(e)

305 - raise e from None

306 -

307 - def _handle_process_request(self, request: RPCProcessRequest):

308 - """Handle RPCProcessRequest by adding it to the LLMEngine."""

309 - request_id = request.request_id

310 -

311 - if self._errored_with is not None:

312 - rpc_err = RPCError(request_id=request_id,

313 - is_engine_errored=True,

314 - exception=ENGINE_DEAD_ERROR(self._errored_with))

315 - self._send_outputs(rpc_err)

316 -

317 - try:

318 - self.engine.add_request(request_id=request_id,

319 - prompt=request.prompt,

320 - params=request.params,

321 - lora_request=request.lora_request,

322 - trace_headers=request.trace_headers,

323 - priority=request.priority)

324 -

325 - if self.log_requests:

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

327 -

328 - except Exception as e:

329 - # We do not set self._errored = True here, since the error

330 - # is due to an issue adding this request to the engine,

331 - # rather than an issue with the engine itself.

332 - logger.debug("Failed to add request %s to engine. %s",

333 - request.request_id, e)

334 - is_errored = self._errored_with is not None

335 - rpc_err = RPCError(request_id=request_id,

336 - is_engine_errored=is_errored,

337 - exception=e)

338 - self._send_outputs(rpc_err)

339 -

340 - # Remove request from the engine.

341 - self.engine.abort_request(request_id)

342 -

343 - def _handle_abort_request(self, request: RPCAbortRequest):

344 - self.engine.abort_request(request.request_id)

345 - if self.log_requests:

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

347 -

348 - def _handle_load_adapter_request(self, request: RPCLoadAdapterRequest):

349 - try:

350 - lora_loaded = self.engine.add_lora(request.lora_request)

351 - except BaseException as e:

352 - # Send back an error if the adater fails to load

353 - rpc_err = RPCError(request_id=request.request_id,

354 - is_engine_errored=False,

355 - exception=e)

356 - self._send_outputs(rpc_err)

357 - return

358 - # Otherwise, send back the successful load message

359 - self._send_outputs(

360 - RPCAdapterLoadedResponse(request_id=request.request_id,

361 - lora_loaded=lora_loaded))

362 -

363 - def _handle_is_sleeping_request(self, request: RPCIsSleepingRequest):

364 - is_sleeping = self.is_sleeping()

365 - self._send_outputs(

366 - RPCIsSleepingResponse(request_id=request.request_id,

367 - is_sleeping=is_sleeping))

368 -

369 - def _health_check(self):

370 - # Send unhealthy if engine has already errored

371 - if self._errored_with is not None:

372 - self._send_unhealthy(self._errored_with)

373 - try:

374 - self.engine.check_health()

375 - self._send_healthy()

376 - except Exception as e:

377 - self._set_errored(e)

378 - self._send_unhealthy(e)

379 -

380 - def _send_outputs(self, outputs: REQUEST_OUTPUTS_T):

381 - """Send outputs back to the engine client. These can be:

382 - - Exceptions

383 - - A list of generation outputs

384 - - A response from loading a lora adapter

385 - """

386 - if outputs:

387 - try:

388 - from ray.exceptions import RayTaskError

389 -

390 - # RayTaskError might not pickelable here. We need to unpack the

391 - # underlying exception as the real exception in the output.

392 - if (isinstance(outputs, RPCError)

393 - and isinstance(outputs.exception, RayTaskError)):

394 - outputs.exception = outputs.exception.cause

395 - except ImportError:

396 - pass

397 -

398 - output_bytes = pickle.dumps(outputs)

399 - self.output_socket.send_multipart((output_bytes, ), copy=False)

400 -

401 - def _send_healthy(self):

402 - """Send HEALTHY message to RPCClient."""

403 - if not self.heartbeat_socket.closed:

404 - self.heartbeat_socket.send_multipart(HEALTHY_RESPONSE, copy=False)

405 -

406 - def _send_unhealthy(self, error: BaseException):

407 - """Send UNHEALTHY message to RPCClient."""

408 - if not self.heartbeat_socket.closed:

409 - error_bytes = pickle.dumps(error)

410 - self.heartbeat_socket.send_multipart((error_bytes, ), copy=False)

411 -

412 - def _async_socket_engine_callback(self,

413 - request_outputs: REQUEST_OUTPUTS_T):

414 - """Callback used by engine to make socket handling async with GPU."""

415 - self._send_outputs(request_outputs)

416 - self.handle_new_input()

417 -

418 - def _set_errored(self, e: BaseException):

419 - """Log and set errored status if this is the first issue."""

420 - if self._errored_with is None:

421 - self._errored_with = e

422 -

423 - def start_profile(self) -> None:

424 - self.engine.start_profile()

425 -

426 - def stop_profile(self) -> None:

427 - self.engine.stop_profile()

428 -

429 - def reset_mm_cache(self) -> bool:

430 - return self.engine.reset_mm_cache()

431 -

432 - def reset_prefix_cache(self) -> bool:

433 - return self.engine.reset_prefix_cache()

434 -

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

436 - self.engine.sleep(level)

437 -

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

439 - self.engine.wake_up(tags)

440 -

441 - def is_sleeping(self) -> bool:

442 - return self.engine.is_sleeping()

443 -

444 -

445 -def signal_handler(*_) -> None:

446 - raise KeyboardInterrupt("MQLLMEngine terminated")

447 -

448 -

449 -def run_mp_engine(vllm_config: VllmConfig, usage_context: UsageContext,

450 - ipc_path: str, disable_log_stats: bool,

451 - enable_log_requests: bool, engine_alive):

452 - try:

453 - # Ensure we can serialize transformer config before spawning

454 - maybe_register_config_serialize_by_value()

455 -

456 - engine = MQLLMEngine.from_vllm_config(

457 - vllm_config=vllm_config,

458 - usage_context=usage_context,

459 - disable_log_stats=disable_log_stats,

460 - enable_log_requests=enable_log_requests,

461 - ipc_path=ipc_path)

462 -

463 - signal.signal(signal.SIGTERM, signal_handler)

464 -

465 - engine.start()

466 -

467 - except BaseException as e:

468 - logger.exception(e)

469 - engine_alive.value = False

470 - raise e from None

vllm/entrypoints/launcher.py

-2

@@ -12,7 +12,6 @@ from fastapi import FastAPI, Request, Response

13 from vllm import envs

14 from vllm.engine.async_llm_engine import AsyncEngineDeadError

15 -from vllm.engine.multiprocessing import MQEngineDeadError

15 from vllm.engine.protocol import EngineClient

16 from vllm.entrypoints.constants import (H11_MAX_HEADER_COUNT_DEFAULT,

17 H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT)

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

155

156 @app.exception_handler(RuntimeError)

157 @app.exception_handler(AsyncEngineDeadError)

159 - @app.exception_handler(MQEngineDeadError)

158 @app.exception_handler(EngineDeadError)

159 @app.exception_handler(EngineGenerateError)

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

vllm/entrypoints/openai/api_server.py

+2 -100

@@ -2,7 +2,6 @@

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

4 import asyncio

5 -import atexit

5 import gc

6 import importlib

7 import inspect

@@ -17,7 +16,6 @@ import uuid

16 from argparse import Namespace

17 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable

18 from contextlib import asynccontextmanager

20 -from functools import partial

19 from http import HTTPStatus

20 from typing import Annotated, Any, Callable, Optional

@@ -42,8 +40,6 @@ import vllm.envs as envs

40 from vllm.config import VllmConfig

41 from vllm.engine.arg_utils import AsyncEngineArgs

42 from vllm.engine.async_llm_engine import AsyncLLMEngine # type: ignore

45 -from vllm.engine.multiprocessing.client import MQLLMEngineClient

46 -from vllm.engine.multiprocessing.engine import run_mp_engine

43 from vllm.engine.protocol import EngineClient

44 from vllm.entrypoints.chat_utils import (load_chat_template,

45 resolve_hf_chat_template,

@@ -102,13 +98,10 @@ from vllm.entrypoints.utils import (cli_env_setup, load_aware_call,

98 log_non_default_args, with_cancellation)

99 from vllm.logger import init_logger

100 from vllm.reasoning import ReasoningParserManager

105 -from vllm.transformers_utils.config import (

106 - maybe_register_config_serialize_by_value)

101 from vllm.transformers_utils.tokenizer import MistralTokenizer

102 from vllm.usage.usage_lib import UsageContext

103 from vllm.utils import (Device, FlexibleArgumentParser, decorate_logs,

110 - get_open_zmq_ipc_path, is_valid_ipv6_address,

111 - set_ulimit)

104 + is_valid_ipv6_address, set_ulimit)

105 from vllm.v1.metrics.prometheus import get_prometheus_registry

106 from vllm.version import __version__ as VLLM_VERSION

107

@@ -237,8 +230,7 @@ async def build_async_engine_client_from_engine_args(

230 async_llm.shutdown()

231

232 # V0 AsyncLLM.

240 - elif (MQLLMEngineClient.is_unsupported_config(vllm_config)

241 - or disable_frontend_multiprocessing):

233 + else:

234

235 engine_client: Optional[EngineClient] = None

236 try:

@@ -252,96 +244,6 @@ async def build_async_engine_client_from_engine_args(

244 if engine_client and hasattr(engine_client, "shutdown"):

245 engine_client.shutdown()

246

255 - # V0MQLLMEngine.

256 - else:

257 - if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:

258 - # Make TemporaryDirectory for prometheus multiprocessing

259 - # Note: global TemporaryDirectory will be automatically

260 - # cleaned up upon exit.

261 - global prometheus_multiproc_dir

262 - prometheus_multiproc_dir = tempfile.TemporaryDirectory()

263 - os.environ[

264 - "PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name

265 - else:

266 - logger.warning(

267 - "Found PROMETHEUS_MULTIPROC_DIR was set by user. "

268 - "This directory must be wiped between vLLM runs or "

269 - "you will find inaccurate metrics. Unset the variable "

270 - "and vLLM will properly handle cleanup.")

271 -

272 - # Select random path for IPC.

273 - ipc_path = get_open_zmq_ipc_path()

274 - logger.debug("Multiprocessing frontend to use %s for IPC Path.",

275 - ipc_path)

276 -

277 - # Start RPCServer in separate process (holds the LLMEngine).

278 - # the current process might have CUDA context,

279 - # so we need to spawn a new process

280 - context = multiprocessing.get_context("spawn")

281 -

282 - # Ensure we can serialize transformer config before spawning

283 - maybe_register_config_serialize_by_value()

284 -

285 - # The Process can raise an exception during startup, which may

286 - # not actually result in an exitcode being reported. As a result

287 - # we use a shared variable to communicate the information.

288 - engine_alive = multiprocessing.Value('b', True, lock=False)

289 - engine_process = context.Process(

290 - target=run_mp_engine,

291 - args=(vllm_config, UsageContext.OPENAI_API_SERVER, ipc_path,

292 - engine_args.disable_log_stats,

293 - engine_args.enable_log_requests, engine_alive))

294 - engine_process.start()

295 - engine_pid = engine_process.pid

296 - assert engine_pid is not None, "Engine process failed to start."

297 - logger.info("Started engine process with PID %d", engine_pid)

298 -

299 - def _cleanup_ipc_path():

300 - socket_path = ipc_path.replace("ipc://", "")

301 - if os.path.exists(socket_path):

302 - os.remove(socket_path)

303 -

304 - # Ensure we clean up the local IPC socket file on exit.

305 - atexit.register(_cleanup_ipc_path)

306 -

307 - # Build RPCClient, which conforms to EngineClient Protocol.

308 - build_client = partial(MQLLMEngineClient, ipc_path, vllm_config,

309 - engine_pid)

310 - mq_engine_client = await asyncio.get_running_loop().run_in_executor(

311 - None, build_client)

312 - try:

313 - while True:

314 - try:

315 - await mq_engine_client.setup()

316 - break

317 - except TimeoutError:

318 - if (not engine_process.is_alive()

319 - or not engine_alive.value):

320 - raise RuntimeError(

321 - "Engine process failed to start. See stack "

322 - "trace for the root cause.") from None

323 -

324 - yield mq_engine_client # type: ignore[misc]

325 - finally:

326 - # Ensure rpc server process was terminated

327 - engine_process.terminate()

328 -

329 - # Close all open connections to the backend

330 - mq_engine_client.close()

331 -

332 - # Wait for engine process to join

333 - engine_process.join(4)

334 - if engine_process.exitcode is None:

335 - # Kill if taking longer than 5 seconds to stop

336 - engine_process.kill()

337 -

338 - # Lazy import for prometheus multiprocessing.

339 - # We need to set PROMETHEUS_MULTIPROC_DIR environment variable

340 - # before prometheus_client is imported.

341 - # See https://prometheus.github.io/client_python/multiprocess/

342 - from prometheus_client import multiprocess

343 - multiprocess.mark_process_dead(engine_process.pid)

344 -

247

248 async def validate_json_request(raw_request: Request):

249 content_type = raw_request.headers.get("content-type", "").lower()

vllm/platforms/rocm.py

+1 -1

@@ -191,7 +191,7 @@ class RocmPlatform(Platform):

191 kv_cache_dtype, block_size, use_v1, use_mla,

192 has_sink) -> str:

193 if use_mla:

194 - from vllm.attention.backends.rocm_aiter_mla import (

194 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import (

195 is_aiter_mla_enabled)

196

197 if selected_backend is None: