74bc397b0a · tduncan/Nvidia-Optimization-Stack-VLLM

tests/entrypoints/openai/test_sleep.py

+7

@@ -28,5 +28,12 @@ def test_sleep_mode():

28 response = requests.post(remote_server.url_for("/sleep"),

29 data={"level": "1"})

30 assert response.status_code == 200

31 + response = requests.get(remote_server.url_for("/is_sleeping"))

32 + assert response.status_code == 200

33 + assert response.json().get("is_sleeping") is True

34 +

35 response = requests.post(remote_server.url_for("/wake_up"))

36 assert response.status_code == 200

37 + response = requests.get(remote_server.url_for("/is_sleeping"))

38 + assert response.status_code == 200

39 + assert response.json().get("is_sleeping") is False

vllm/engine/async_llm_engine.py

+3

@@ -1225,6 +1225,9 @@ class AsyncLLMEngine(EngineClient):

1225 async def wake_up(self) -> None:

1226 self.engine.wake_up()

1227

1228 + async def is_sleeping(self) -> bool:

1229 + return self.engine.is_sleeping()

1230 +

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

1232 self.engine.add_lora(lora_request)

1233

vllm/engine/llm_engine.py

+3

@@ -1948,6 +1948,9 @@ class LLMEngine:

1948 "Sleep mode is not enabled in the model config")

1949 self.model_executor.wake_up()

1950

1951 + def is_sleeping(self) -> bool:

1952 + return self.model_executor.is_sleeping

1953 +

1954 def check_health(self) -> None:

1955 if self.tokenizer:

1956 self.tokenizer.check_health()

vllm/engine/multiprocessing/__init__.py

+14 -2

@@ -136,6 +136,18 @@ class RPCWakeUpRequest(Enum):

136 WAKE_UP = 1

137

138

139 +@dataclass

140 +class RPCIsSleepingRequest:

141 + # Set the default value of request_id to a new UUID

142 + request_id: str = field(default_factory=lambda: str(uuid.uuid4()))

143 +

144 +

145 +@dataclass

146 +class RPCIsSleepingResponse:

147 + request_id: str

148 + is_sleeping: bool

149 +

150 +

151 @dataclass

152 class RPCLoadAdapterRequest:

153 lora_request: LoRARequest

@@ -151,10 +163,10 @@ class RPCAdapterLoadedResponse:

163 RPC_REQUEST_T = Union[RPCProcessRequest, RPCAbortRequest, RPCStartupRequest,

164 RPCUProfileRequest, RPCLoadAdapterRequest,

165 RPCResetPrefixCacheRequest, RPCSleepRequest,

154 - RPCWakeUpRequest]

166 + RPCWakeUpRequest, RPCIsSleepingRequest]

167

168 REQUEST_OUTPUTS_T = Union[List[RequestOutput], RPCAdapterLoadedResponse,

157 - RPCError]

169 + RPCIsSleepingResponse, RPCError]

170

171

172 def ENGINE_DEAD_ERROR(

vllm/engine/multiprocessing/client.py

+25 -2

@@ -27,6 +27,8 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

27 IPC_OUTPUT_EXT, RPC_REQUEST_T,

28 VLLM_RPC_SUCCESS_STR, RPCAbortRequest,

29 RPCAdapterLoadedResponse, RPCError,

30 + RPCIsSleepingRequest,

31 + RPCIsSleepingResponse,

32 RPCLoadAdapterRequest,

33 RPCProcessRequest,

34 RPCResetPrefixCacheRequest,

@@ -246,7 +248,9 @@ class MQLLMEngineClient(EngineClient):

248 if queue is not None:

249 queue.put_nowait(exception)

250 # Put each output into the appropriate queue.

249 - elif isinstance(request_outputs, RPCAdapterLoadedResponse):

251 + elif isinstance(

252 + request_outputs,

253 + (RPCAdapterLoadedResponse, RPCIsSleepingResponse)):

254 self._add_output(request_outputs)

255 else:

256 for request_output in request_outputs:

@@ -256,7 +260,8 @@ class MQLLMEngineClient(EngineClient):

260 logger.debug("Shutting down MQLLMEngineClient output handler.")

261

262 def _add_output(self, request_output: Union[RequestOutput,

259 - RPCAdapterLoadedResponse]):

263 + RPCAdapterLoadedResponse,

264 + RPCIsSleepingResponse]):

265 queue = self.output_queues.get(request_output.request_id)

266 if queue is not None:

267 queue.put_nowait(request_output)

@@ -696,6 +701,24 @@ class MQLLMEngineClient(EngineClient):

701 return await self._send_one_way_rpc_request(

702 request=RPCWakeUpRequest.WAKE_UP, socket=self.input_socket)

703

704 + async def is_sleeping(self) -> bool:

705 + """Check whether the engine is sleeping"""

706 + request = RPCIsSleepingRequest()

707 +

708 + queue: asyncio.Queue[Union[BaseException,

709 + RPCIsSleepingResponse]] = asyncio.Queue()

710 + self.output_queues[request.request_id] = queue

711 +

712 + request_bytes = pickle.dumps(request)

713 + await self.input_socket.send_multipart((request_bytes, ), copy=False)

714 +

715 + request_output = await queue.get()

716 + self.output_queues.pop(request.request_id)

717 +

718 + if isinstance(request_output, BaseException):

719 + raise request_output

720 + return request_output.is_sleeping

721 +

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

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

724 # Uses the same I/O as generate requests

vllm/engine/multiprocessing/engine.py

+13

@@ -18,6 +18,8 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

18 IPC_OUTPUT_EXT, REQUEST_OUTPUTS_T,

19 VLLM_RPC_SUCCESS_STR, RPCAbortRequest,

20 RPCAdapterLoadedResponse, RPCError,

21 + RPCIsSleepingRequest,

22 + RPCIsSleepingResponse,

23 RPCLoadAdapterRequest,

24 RPCProcessRequest,

25 RPCResetPrefixCacheRequest,

@@ -271,6 +273,8 @@ class MQLLMEngine:

273 self.sleep(request.value)

274 elif isinstance(request, RPCWakeUpRequest):

275 self.wake_up()

276 + elif isinstance(request, RPCIsSleepingRequest):

277 + self._handle_is_sleeping_request(request)

278 else:

279 raise ValueError("Unknown RPCRequest Type: "

280 f"{type(request)}")

@@ -337,6 +341,12 @@ class MQLLMEngine:

341 self._send_outputs(

342 RPCAdapterLoadedResponse(request_id=request.request_id))

343

344 + def _handle_is_sleeping_request(self, request: RPCIsSleepingRequest):

345 + is_sleeping = self.is_sleeping()

346 + self._send_outputs(

347 + RPCIsSleepingResponse(request_id=request.request_id,

348 + is_sleeping=is_sleeping))

349 +

350 def _health_check(self):

351 # Send unhealthy if engine has already errored

352 if self._errored_with is not None:

@@ -406,6 +416,9 @@ class MQLLMEngine:

416 def wake_up(self) -> None:

417 self.engine.wake_up()

418

419 + def is_sleeping(self) -> bool:

420 + return self.engine.is_sleeping()

421 +

422

423 def signal_handler(*_) -> None:

424 raise KeyboardInterrupt("MQLLMEngine terminated")

vllm/engine/protocol.py

+5

@@ -288,6 +288,11 @@ class EngineClient(ABC):

288 """Wake up the engine"""

289 ...

290

291 + @abstractmethod

292 + async def is_sleeping(self) -> bool:

293 + """Check whether the engine is sleeping"""

294 + ...

295 +

296 @abstractmethod

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

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

vllm/entrypoints/openai/api_server.py

+6

@@ -694,6 +694,12 @@ if envs.VLLM_SERVER_DEV_MODE:

694 # is sent but does not finish yet when we return a response.

695 return Response(status_code=200)

696

697 + @router.get("/is_sleeping")

698 + async def is_sleeping(raw_request: Request):

699 + logger.info("check whether the engine is sleeping")

700 + is_sleeping = await engine_client(raw_request).is_sleeping()

701 + return JSONResponse(content={"is_sleeping": is_sleeping})

702 +

703

704 @router.post("/invocations", dependencies=[Depends(validate_json_request)])

705 async def invocations(raw_request: Request):

vllm/v1/engine/async_llm.py

+3

@@ -407,6 +407,9 @@ class AsyncLLM(EngineClient):

407 async def wake_up(self) -> None:

408 await self.engine_core.wake_up_async()

409

410 + async def is_sleeping(self) -> bool:

411 + return await self.engine_core.is_sleeping_async()

412 +

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

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

415 return await self.engine_core.add_lora_async(lora_request)

vllm/v1/engine/core.py

+3

@@ -253,6 +253,9 @@ class EngineCore:

253 def wake_up(self):

254 self.model_executor.wake_up()

255

256 + def is_sleeping(self) -> bool:

257 + return self.model_executor.is_sleeping

258 +

259 def execute_dummy_batch(self):

260 self.model_executor.collective_rpc("execute_dummy_batch")

261

vllm/v1/engine/core_client.py

+15

@@ -89,6 +89,9 @@ class EngineCoreClient(ABC):

89 def wake_up(self) -> None:

90 raise NotImplementedError

91

92 + def is_sleeping(self) -> bool:

93 + raise NotImplementedError

94 +

95 def execute_dummy_batch(self) -> None:

96 raise NotImplementedError

97

@@ -128,6 +131,9 @@ class EngineCoreClient(ABC):

131 async def wake_up_async(self) -> None:

132 raise NotImplementedError

133

134 + async def is_sleeping_async(self) -> bool:

135 + raise NotImplementedError

136 +

137 async def abort_requests_async(self, request_ids: list[str]) -> None:

138 raise NotImplementedError

139

@@ -182,6 +188,9 @@ class InprocClient(EngineCoreClient):

188 def wake_up(self) -> None:

189 self.engine_core.wake_up()

190

191 + def is_sleeping(self) -> bool:

192 + return self.engine_core.is_sleeping()

193 +

194 def execute_dummy_batch(self) -> None:

195 self.engine_core.execute_dummy_batch()

196

@@ -433,6 +442,9 @@ class SyncMPClient(MPClient):

442 def wake_up(self) -> None:

443 self._call_utility("wake_up")

444

445 + def is_sleeping(self) -> bool:

446 + return self._call_utility("is_sleeping")

447 +

448 def execute_dummy_batch(self) -> None:

449 self._call_utility("execute_dummy_batch")

450

@@ -523,6 +535,9 @@ class AsyncMPClient(MPClient):

535 async def wake_up_async(self) -> None:

536 await self._call_utility_async("wake_up")

537

538 + async def is_sleeping_async(self) -> bool:

539 + return await self._call_utility_async("is_sleeping")

540 +

541 async def execute_dummy_batch_async(self) -> None:

542 await self._call_utility_async("execute_dummy_batch")

543

vllm/v1/engine/llm_engine.py

+3

@@ -235,6 +235,9 @@ class LLMEngine:

235 def wake_up(self):

236 self.engine_core.wake_up()

237

238 + def is_sleeping(self) -> bool:

239 + return self.engine_core.is_sleeping()

240 +

241 def get_tokenizer_group(

242 self,

243 group_type: type[_G] = BaseTokenizerGroup,

[Core] Expose API endpoint `/is_sleeping` (#14312)