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

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 15, 22:15 UTC

[V1][Core] Generic mechanism for handling engine utility (#13060)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Feb 19, 2025 at 01:09 UTC caf7ff4456d5c03b31cad89d02d03eb8d158e815

5 files changed +197 -56

tests/lora/test_add_lora.py

+1 -1

@@ -41,7 +41,7 @@ def download_and_prepare_lora_module():

41 ]

42 for tokenizer_file in tokenizer_files:

43 del_path = Path(LORA_MODULE_DOWNLOAD_PATH) / tokenizer_file

44 - del_path.unlink()

44 + del_path.unlink(missing_ok=True)

47 @pytest.fixture(autouse=True)

tests/v1/engine/test_engine_core_client.py

+49 -8

@@ -3,7 +3,8 @@

3 import asyncio

4 import time

5 import uuid

6 -from typing import Dict, List

6 +from contextlib import ExitStack

7 +from typing import Dict, List, Optional

9 import pytest

10 from transformers import AutoTokenizer

@@ -14,7 +15,9 @@ from vllm.engine.arg_utils import EngineArgs

15 from vllm.platforms import current_platform

16 from vllm.usage.usage_lib import UsageContext

17 from vllm.v1.engine import EngineCoreRequest

17 -from vllm.v1.engine.core_client import EngineCoreClient

18 +from vllm.v1.engine.core import EngineCore

19 +from vllm.v1.engine.core_client import (AsyncMPClient, EngineCoreClient,

20 + SyncMPClient)

21 from vllm.v1.executor.abstract import Executor

23 if not current_platform.is_cuda():

@@ -63,7 +66,7 @@ def loop_until_done(client: EngineCoreClient, outputs: Dict):

66 async def loop_until_done_async(client: EngineCoreClient, outputs: Dict):

68 while True:

66 - engine_core_outputs = await client.get_output_async().outputs

69 + engine_core_outputs = (await client.get_output_async()).outputs

71 if len(engine_core_outputs) == 0:

72 break

@@ -78,6 +81,14 @@ async def loop_until_done_async(client: EngineCoreClient, outputs: Dict):

81 break

84 +# Dummy utility function to monkey-patch into engine core.

85 +def echo(self, msg: str, err_msg: Optional[str] = None) -> str:

86 + print(f"echo util function called: {msg}, {err_msg}")

87 + if err_msg is not None:

88 + raise ValueError(err_msg)

89 + return msg

90 +

91 +

92 @fork_new_process_for_each_test

93 @pytest.mark.parametrize("multiprocessing_mode", [True, False])

94 def test_engine_core_client(monkeypatch, multiprocessing_mode: bool):

@@ -85,7 +96,10 @@ def test_engine_core_client(monkeypatch, multiprocessing_mode: bool):

96 with monkeypatch.context() as m:

97 m.setenv("VLLM_USE_V1", "1")

88 - engine_args = EngineArgs(model=MODEL_NAME, compilation_config=3)

99 + # Monkey-patch core engine utility function to test.

100 + m.setattr(EngineCore, "echo", echo, raising=False)

101 +

102 + engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)

103 vllm_config = engine_args.create_engine_config(

104 UsageContext.UNKNOWN_CONTEXT)

105 executor_class = Executor.get_class(vllm_config)

@@ -147,15 +161,30 @@ def test_engine_core_client(monkeypatch, multiprocessing_mode: bool):

161

162 client.abort_requests([request.request_id])

163

164 + if multiprocessing_mode:

165 + """Utility method invocation"""

166

151 -@fork_new_process_for_each_test

152 -@pytest.mark.asyncio

167 + core_client: SyncMPClient = client

168 +

169 + result = core_client._call_utility("echo", "testarg")

170 + assert result == "testarg"

171 +

172 + with pytest.raises(Exception) as e_info:

173 + core_client._call_utility("echo", None, "help!")

174 +

175 + assert str(e_info.value) == "Call to echo method failed: help!"

176 +

177 +

178 +@pytest.mark.asyncio(loop_scope="function")

179 async def test_engine_core_client_asyncio(monkeypatch):

180

155 - with monkeypatch.context() as m:

181 + with monkeypatch.context() as m, ExitStack() as after:

182 m.setenv("VLLM_USE_V1", "1")

183

158 - engine_args = EngineArgs(model=MODEL_NAME)

184 + # Monkey-patch core engine utility function to test.

185 + m.setattr(EngineCore, "echo", echo, raising=False)

186 +

187 + engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)

188 vllm_config = engine_args.create_engine_config(

189 usage_context=UsageContext.UNKNOWN_CONTEXT)

190 executor_class = Executor.get_class(vllm_config)

@@ -166,6 +195,7 @@ async def test_engine_core_client_asyncio(monkeypatch):

195 executor_class=executor_class,

196 log_stats=True,

197 )

198 + after.callback(client.shutdown)

199

200 MAX_TOKENS = 20

201 params = SamplingParams(max_tokens=MAX_TOKENS)

@@ -204,3 +234,14 @@ async def test_engine_core_client_asyncio(monkeypatch):

234 else:

235 assert len(outputs[req_id]) == MAX_TOKENS, (

236 f"{len(outputs[req_id])=}, {MAX_TOKENS=}")

237 + """Utility method invocation"""

238 +

239 + core_client: AsyncMPClient = client

240 +

241 + result = await core_client._call_utility_async("echo", "testarg")

242 + assert result == "testarg"

243 +

244 + with pytest.raises(Exception) as e_info:

245 + await core_client._call_utility_async("echo", None, "help!")

246 +

247 + assert str(e_info.value) == "Call to echo method failed: help!"

vllm/v1/engine/__init__.py

+18 -6

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

3 import enum

4 import time

5 -from typing import List, Optional, Union

5 +from typing import Any, List, Optional, Union

7 import msgspec

@@ -106,6 +106,18 @@ class EngineCoreOutput(

106 return self.finish_reason is not None

107

108

109 +class UtilityOutput(

110 + msgspec.Struct,

111 + array_like=True, # type: ignore[call-arg]

112 + gc=False): # type: ignore[call-arg]

113 +

114 + call_id: int

115 +

116 + # Non-None implies the call failed, result should be None.

117 + failure_message: Optional[str] = None

118 + result: Any = None

119 +

120 +

121 class EngineCoreOutputs(

122 msgspec.Struct,

123 array_like=True, # type: ignore[call-arg]

@@ -116,10 +128,12 @@ class EngineCoreOutputs(

128 # e.g. columnwise layout

129

130 # [num_reqs]

119 - outputs: List[EngineCoreOutput]

120 - scheduler_stats: Optional[SchedulerStats]

131 + outputs: List[EngineCoreOutput] = []

132 + scheduler_stats: Optional[SchedulerStats] = None

133 timestamp: float = 0.0

134

135 + utility_output: Optional[UtilityOutput] = None

136 +

137 def __post_init__(self):

138 if self.timestamp == 0.0:

139 self.timestamp = time.monotonic()

@@ -132,6 +146,4 @@ class EngineCoreRequestType(enum.Enum):

146 """

147 ADD = b'\x00'

148 ABORT = b'\x01'

135 - PROFILE = b'\x02'

136 - RESET_PREFIX_CACHE = b'\x03'

137 - ADD_LORA = b'\x04'

149 + UTILITY = b'\x02'

vllm/v1/engine/core.py

+33 -16

@@ -5,9 +5,11 @@ import signal

5 import threading

6 import time

7 from concurrent.futures import Future

8 +from inspect import isclass, signature

9 from multiprocessing.connection import Connection

10 from typing import Any, List, Optional, Tuple, Type

12 +import msgspec

13 import psutil

14 import zmq

15 import zmq.asyncio

@@ -21,7 +23,7 @@ from vllm.utils import get_exception_traceback, zmq_socket_ctx

23 from vllm.v1.core.kv_cache_utils import get_kv_cache_configs

24 from vllm.v1.core.scheduler import Scheduler, SchedulerOutput

25 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

24 - EngineCoreRequestType)

26 + EngineCoreRequestType, UtilityOutput)

27 from vllm.v1.engine.mm_input_cache import MMInputCacheServer

28 from vllm.v1.executor.abstract import Executor

29 from vllm.v1.outputs import ModelRunnerOutput

@@ -330,19 +332,39 @@ class EngineCoreProc(EngineCore):

332 self.add_request(request)

333 elif request_type == EngineCoreRequestType.ABORT:

334 self.abort_requests(request)

333 - elif request_type == EngineCoreRequestType.RESET_PREFIX_CACHE:

334 - self.reset_prefix_cache()

335 - elif request_type == EngineCoreRequestType.PROFILE:

336 - self.model_executor.profile(request)

337 - elif request_type == EngineCoreRequestType.ADD_LORA:

338 - self.model_executor.add_lora(request)

335 + elif request_type == EngineCoreRequestType.UTILITY:

336 + call_id, method_name, args = request

337 + output = UtilityOutput(call_id)

338 + try:

339 + method = getattr(self, method_name)

340 + output.result = method(

341 + *self._convert_msgspec_args(method, args))

342 + except BaseException as e:

343 + logger.exception("Invocation of %s method failed", method_name)

344 + output.failure_message = (f"Call to {method_name} method"

345 + f" failed: {str(e)}")

346 + self.output_queue.put_nowait(

347 + EngineCoreOutputs(utility_output=output))

348 +

349 + @staticmethod

350 + def _convert_msgspec_args(method, args):

351 + """If a provided arg type doesn't match corresponding target method

352 + arg type, try converting to msgspec object."""

353 + if not args:

354 + return args

355 + arg_types = signature(method).parameters.values()

356 + assert len(args) <= len(arg_types)

357 + return tuple(

358 + msgspec.convert(v, type=p.annotation) if isclass(p.annotation)

359 + and issubclass(p.annotation, msgspec.Struct)

360 + and not isinstance(v, p.annotation) else v

361 + for v, p in zip(args, arg_types))

362

363 def process_input_socket(self, input_path: str):

364 """Input socket IO thread."""

365

366 # Msgpack serialization decoding.

367 add_request_decoder = MsgpackDecoder(EngineCoreRequest)

345 - add_lora_decoder = MsgpackDecoder(LoRARequest)

368 generic_decoder = MsgpackDecoder()

369

370 with zmq_socket_ctx(input_path, zmq.constants.PULL) as socket:

@@ -352,14 +374,9 @@ class EngineCoreProc(EngineCore):

374 request_type = EngineCoreRequestType(bytes(type_frame.buffer))

375

376 # Deserialize the request data.

355 - decoder = None

356 - if request_type == EngineCoreRequestType.ADD:

357 - decoder = add_request_decoder

358 - elif request_type == EngineCoreRequestType.ADD_LORA:

359 - decoder = add_lora_decoder

360 - else:

361 - decoder = generic_decoder

362 -

377 + decoder = add_request_decoder if (

378 + request_type

379 + == EngineCoreRequestType.ADD) else generic_decoder

380 request = decoder.decode(data_frame.buffer)

381

382 # Push to input queue for core busy loop.

vllm/v1/engine/core_client.py

+96 -25

@@ -2,10 +2,14 @@

3 import asyncio

4 import os

5 +import queue

6 import signal

7 +import uuid

8 import weakref

9 from abc import ABC, abstractmethod

8 -from typing import Any, List, Optional, Type

10 +from concurrent.futures import Future

11 +from threading import Thread

12 +from typing import Any, Dict, List, Optional, Type, Union

14 import zmq

15 import zmq.asyncio

@@ -16,7 +20,7 @@ from vllm.lora.request import LoRARequest

20 from vllm.utils import (get_open_zmq_ipc_path, kill_process_tree,

21 make_zmq_socket)

22 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

19 - EngineCoreRequestType)

23 + EngineCoreRequestType, UtilityOutput)

24 from vllm.v1.engine.core import EngineCore, EngineCoreProc

25 from vllm.v1.executor.abstract import Executor

26 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

@@ -24,6 +28,8 @@ from vllm.v1.utils import BackgroundProcHandle

29 logger = init_logger(__name__)

31 +AnyFuture = Union[asyncio.Future[Any], Future[Any]]

32 +

34 class EngineCoreClient(ABC):

35 """

@@ -204,6 +210,8 @@ class MPClient(EngineCoreClient):

210 "log_stats": log_stats,

211 })

212

213 + self.utility_results: Dict[int, AnyFuture] = {}

214 +

215 def shutdown(self):

216 """Clean up background resources."""

217 if hasattr(self, "proc_handle"):

@@ -212,6 +220,16 @@ class MPClient(EngineCoreClient):

220 self._finalizer()

221

222

223 +def _process_utility_output(output: UtilityOutput,

224 + utility_results: Dict[int, AnyFuture]):

225 + """Set the result from a utility method in the waiting future"""

226 + future = utility_results.pop(output.call_id)

227 + if output.failure_message is not None:

228 + future.set_exception(Exception(output.failure_message))

229 + else:

230 + future.set_result(output.result)

231 +

232 +

233 class SyncMPClient(MPClient):

234 """Synchronous client for multi-proc EngineCore."""

235

@@ -224,10 +242,30 @@ class SyncMPClient(MPClient):

242 log_stats=log_stats,

243 )

244

227 - def get_output(self) -> EngineCoreOutputs:

245 + self.outputs_queue: queue.Queue[EngineCoreOutputs] = queue.Queue()

246

229 - (frame, ) = self.output_socket.recv_multipart(copy=False)

230 - return self.decoder.decode(frame.buffer)

247 + # Ensure that the outputs socket processing thread does not have

248 + # a ref to the client which prevents gc.

249 + output_socket = self.output_socket

250 + decoder = self.decoder

251 + utility_results = self.utility_results

252 + outputs_queue = self.outputs_queue

253 +

254 + def process_outputs_socket():

255 + while True:

256 + (frame, ) = output_socket.recv_multipart(copy=False)

257 + outputs = decoder.decode(frame.buffer)

258 + if outputs.utility_output:

259 + _process_utility_output(outputs.utility_output,

260 + utility_results)

261 + else:

262 + outputs_queue.put_nowait(outputs)

263 +

264 + # Process outputs from engine in separate thread.

265 + Thread(target=process_outputs_socket, daemon=True).start()

266 +

267 + def get_output(self) -> EngineCoreOutputs:

268 + return self.outputs_queue.get()

269

270 def _send_input(self, request_type: EngineCoreRequestType,

271 request: Any) -> None:

@@ -236,6 +274,16 @@ class SyncMPClient(MPClient):

274 msg = (request_type.value, self.encoder.encode(request))

275 self.input_socket.send_multipart(msg, copy=False)

276

277 + def _call_utility(self, method: str, *args) -> Any:

278 + call_id = uuid.uuid1().int >> 64

279 + future: Future[Any] = Future()

280 + self.utility_results[call_id] = future

281 +

282 + self._send_input(EngineCoreRequestType.UTILITY,

283 + (call_id, method, args))

284 +

285 + return future.result()

286 +

287 def add_request(self, request: EngineCoreRequest) -> None:

288 # NOTE: text prompt is not needed in the core engine as it has been

289 # tokenized.

@@ -247,13 +295,13 @@ class SyncMPClient(MPClient):

295 self._send_input(EngineCoreRequestType.ABORT, request_ids)

296

297 def profile(self, is_start: bool = True) -> None:

250 - self._send_input(EngineCoreRequestType.PROFILE, is_start)

298 + self._call_utility("profile", is_start)

299

300 def reset_prefix_cache(self) -> None:

253 - self._send_input(EngineCoreRequestType.RESET_PREFIX_CACHE, None)

301 + self._call_utility("reset_prefix_cache")

302

303 def add_lora(self, lora_request: LoRARequest) -> None:

256 - self._send_input(EngineCoreRequestType.ADD_LORA, lora_request)

304 + self._call_utility("add_lora", lora_request)

305

306

307 class AsyncMPClient(MPClient):

@@ -268,24 +316,35 @@ class AsyncMPClient(MPClient):

316 log_stats=log_stats,

317 )

318

271 - self.outputs_queue: Optional[asyncio.Queue[bytes]] = None

319 + self.outputs_queue: Optional[asyncio.Queue[EngineCoreOutputs]] = None

320 self.queue_task: Optional[asyncio.Task] = None

321

322 + async def _start_output_queue_task(self):

323 + # Perform IO in separate task to parallelize as much as possible.

324 + # Avoid task having direct reference back to the client.

325 + self.outputs_queue = asyncio.Queue()

326 + output_socket = self.output_socket

327 + decoder = self.decoder

328 + utility_results = self.utility_results

329 + outputs_queue = self.outputs_queue

330 +

331 + async def process_outputs_socket():

332 + while True:

333 + (frame, ) = await output_socket.recv_multipart(copy=False)

334 + outputs: EngineCoreOutputs = decoder.decode(frame.buffer)

335 + if outputs.utility_output:

336 + _process_utility_output(outputs.utility_output,

337 + utility_results)

338 + else:

339 + outputs_queue.put_nowait(outputs)

340 +

341 + self.queue_task = asyncio.create_task(process_outputs_socket())

342 +

343 async def get_output_async(self) -> EngineCoreOutputs:

344 if self.outputs_queue is None:

276 - # Perform IO in separate task to parallelize as much as possible

277 - self.outputs_queue = asyncio.Queue()

278 -

279 - async def process_outputs_socket():

280 - assert self.outputs_queue is not None

281 - while True:

282 - (frame, ) = await self.output_socket.recv_multipart(

283 - copy=False)

284 - self.outputs_queue.put_nowait(frame.buffer)

285 -

286 - self.queue_task = asyncio.create_task(process_outputs_socket())

287 -

288 - return self.decoder.decode(await self.outputs_queue.get())

345 + await self._start_output_queue_task()

346 + assert self.outputs_queue is not None

347 + return await self.outputs_queue.get()

348

349 async def _send_input(self, request_type: EngineCoreRequestType,

350 request: Any) -> None:

@@ -293,6 +352,18 @@ class AsyncMPClient(MPClient):

352 msg = (request_type.value, self.encoder.encode(request))

353 await self.input_socket.send_multipart(msg, copy=False)

354

355 + if self.outputs_queue is None:

356 + await self._start_output_queue_task()

357 +

358 + async def _call_utility_async(self, method: str, *args) -> Any:

359 + call_id = uuid.uuid1().int >> 64

360 + future = asyncio.get_running_loop().create_future()

361 + self.utility_results[call_id] = future

362 + await self._send_input(EngineCoreRequestType.UTILITY,

363 + (call_id, method, args))

364 +

365 + return await future

366 +

367 async def add_request_async(self, request: EngineCoreRequest) -> None:

368 # NOTE: text prompt is not needed in the core engine as it has been

369 # tokenized.

@@ -304,10 +375,10 @@ class AsyncMPClient(MPClient):

375 await self._send_input(EngineCoreRequestType.ABORT, request_ids)

376

377 async def profile_async(self, is_start: bool = True) -> None:

307 - await self._send_input(EngineCoreRequestType.PROFILE, is_start)

378 + await self._call_utility_async("profile", is_start)

379

380 async def reset_prefix_cache_async(self) -> None:

310 - await self._send_input(EngineCoreRequestType.RESET_PREFIX_CACHE, None)

381 + await self._call_utility_async("reset_prefix_cache")

382

383 async def add_lora_async(self, lora_request: LoRARequest) -> None:

313 - await self._send_input(EngineCoreRequestType.ADD_LORA, lora_request)

384 + await self._call_utility_async("add_lora", lora_request)