ba81163997 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/basic_correctness/test_cumem.py

+8 -4

@@ -118,14 +118,16 @@ def test_cumem_with_cudagraph():

118

119 @fork_new_process_for_each_test

120 @pytest.mark.parametrize(

121 - "model",

121 + "model, use_v1",

122 [

123 # sleep mode with safetensors

124 - f"{MODEL_WEIGHTS_S3_BUCKET}/Llama-3.2-1B",

124 + (f"{MODEL_WEIGHTS_S3_BUCKET}/Llama-3.2-1B", True),

125 # sleep mode with pytorch checkpoint

126 - "facebook/opt-125m"

126 + ("facebook/opt-125m", False),

127 ])

128 -def test_end_to_end(model):

128 +def test_end_to_end(model: str, use_v1: bool):

129 + import os

130 + os.environ["VLLM_USE_V1"] = "1" if use_v1 else "0"

131 free, total = torch.cuda.mem_get_info()

132 used_bytes_baseline = total - free # in case other process is running

133 load_format = LoadFormat.AUTO

@@ -152,3 +154,5 @@ def test_end_to_end(model):

154

155 # cmp output

156 assert output[0].outputs[0].text == output2[0].outputs[0].text

157 +

158 + del os.environ["VLLM_USE_V1"]

tests/entrypoints/openai/test_sleep.py new

+32

@@ -0,0 +1,32 @@

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

2 +

3 +import requests

4 +

5 +from ...utils import RemoteOpenAIServer

6 +

7 +MODEL_NAME = "meta-llama/Llama-3.2-1B"

8 +

9 +

10 +def test_sleep_mode():

11 + # dtype, max-len etc set so that this can run in CI

12 + args = [

13 + "--dtype",

14 + "bfloat16",

15 + "--max-model-len",

16 + "8192",

17 + "--max-num-seqs",

18 + "128",

19 + "--enable-sleep-mode",

20 + ]

21 +

22 + with RemoteOpenAIServer(MODEL_NAME,

23 + args,

24 + env_dict={

25 + "VLLM_SERVER_DEV_MODE": "1",

26 + "CUDA_VISIBLE_DEVICES": "0"

27 + }) as remote_server:

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

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

30 + assert response.status_code == 200

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

32 + assert response.status_code == 200

vllm/engine/async_llm_engine.py

+6

@@ -1187,6 +1187,12 @@ class AsyncLLMEngine(EngineClient):

1187 async def reset_prefix_cache(self) -> None:

1188 self.engine.reset_prefix_cache()

1189

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

1191 + self.engine.sleep(level)

1192 +

1193 + async def wake_up(self) -> None:

1194 + self.engine.wake_up()

1195 +

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

1197 self.engine.add_lora(lora_request)

1198

vllm/engine/multiprocessing/__init__.py

+11 -1

@@ -127,6 +127,15 @@ class RPCResetPrefixCacheRequest(Enum):

127 RESET_PREFIX_CACHE = 1

128

129

130 +class RPCSleepRequest(Enum):

131 + SLEEP_LEVEL_1 = 1

132 + SLEEP_LEVEL_2 = 2

133 +

134 +

135 +class RPCWakeUpRequest(Enum):

136 + WAKE_UP = 1

137 +

138 +

139 @dataclass

140 class RPCLoadAdapterRequest:

141 lora_request: LoRARequest

@@ -141,7 +150,8 @@ class RPCAdapterLoadedResponse:

150

151 RPC_REQUEST_T = Union[RPCProcessRequest, RPCAbortRequest, RPCStartupRequest,

152 RPCUProfileRequest, RPCLoadAdapterRequest,

144 - RPCResetPrefixCacheRequest]

153 + RPCResetPrefixCacheRequest, RPCSleepRequest,

154 + RPCWakeUpRequest]

155

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

157 RPCError]

vllm/engine/multiprocessing/client.py

+13 -2

@@ -31,8 +31,9 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

31 RPCLoadAdapterRequest,

32 RPCProcessRequest,

33 RPCResetPrefixCacheRequest,

34 - RPCStartupRequest, RPCStartupResponse,

35 - RPCUProfileRequest)

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

@@ -685,6 +686,16 @@ class MQLLMEngineClient(EngineClient):

686 request=RPCResetPrefixCacheRequest.RESET_PREFIX_CACHE,

687 socket=self.input_socket)

688

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

690 + """Sleep the engine for a given level"""

691 + return await self._send_one_way_rpc_request(

692 + request=RPCSleepRequest(level), socket=self.input_socket)

693 +

694 + async def wake_up(self) -> None:

695 + """Wake up the engine"""

696 + return await self._send_one_way_rpc_request(

697 + request=RPCWakeUpRequest.WAKE_UP, socket=self.input_socket)

698 +

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

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

701 # Uses the same I/O as generate requests

vllm/engine/multiprocessing/engine.py

+13 -2

@@ -20,8 +20,9 @@ from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,

20 RPCLoadAdapterRequest,

21 RPCProcessRequest,

22 RPCResetPrefixCacheRequest,

23 - RPCStartupRequest, RPCStartupResponse,

24 - RPCUProfileRequest)

23 + RPCSleepRequest, RPCStartupRequest,

24 + RPCStartupResponse,

25 + RPCUProfileRequest, RPCWakeUpRequest)

26 # yapf: enable

27 from vllm.logger import init_logger

28 from vllm.outputs import RequestOutput

@@ -242,6 +243,10 @@ class MQLLMEngine:

243 self._handle_load_adapter_request(request)

244 elif isinstance(request, RPCResetPrefixCacheRequest):

245 self.reset_prefix_cache()

246 + elif isinstance(request, RPCSleepRequest):

247 + self.sleep(request.value)

248 + elif isinstance(request, RPCWakeUpRequest):

249 + self.wake_up()

250 else:

251 raise ValueError("Unknown RPCRequest Type: "

252 f"{type(request)}")

@@ -369,6 +374,12 @@ class MQLLMEngine:

374 def reset_prefix_cache(self) -> bool:

375 return self.engine.reset_prefix_cache()

376

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

378 + self.engine.sleep(level)

379 +

380 + def wake_up(self) -> None:

381 + self.engine.wake_up()

382 +

383

384 def signal_handler(*_) -> None:

385 raise KeyboardInterrupt("MQLLMEngine terminated")

vllm/engine/protocol.py

+10

@@ -278,6 +278,16 @@ class EngineClient(ABC):

278 """Reset the prefix cache"""

279 ...

280

281 + @abstractmethod

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

283 + """Sleep the engine"""

284 + ...

285 +

286 + @abstractmethod

287 + async def wake_up(self) -> None:

288 + """Wake up the engine"""

289 + ...

290 +

291 @abstractmethod

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

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

vllm/entrypoints/openai/api_server.py

+18

@@ -625,6 +625,24 @@ if envs.VLLM_SERVER_DEV_MODE:

625 await engine_client(raw_request).reset_prefix_cache()

626 return Response(status_code=200)

627

628 + @router.post("/sleep")

629 + async def sleep(raw_request: Request):

630 + # get POST params

631 + level = raw_request.query_params.get("level", "1")

632 + logger.info("sleep the engine with level %s", level)

633 + await engine_client(raw_request).sleep(int(level))

634 + # FIXME: in v0 with frontend multiprocessing, the sleep command

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

636 + return Response(status_code=200)

637 +

638 + @router.post("/wake_up")

639 + async def wake_up(raw_request: Request):

640 + logger.info("wake up the engine")

641 + await engine_client(raw_request).wake_up()

642 + # FIXME: in v0 with frontend multiprocessing, the wake-up command

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

644 + return Response(status_code=200)

645 +

646

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

648 async def invocations(raw_request: Request):

vllm/entrypoints/openai/serving_transcription.py

+1

@@ -295,6 +295,7 @@ class OpenAIServingTranscription(OpenAIServing):

295 # TODO(rob): figure out a way to pipe streaming in.

296 # Non-streaming response.

297 try:

298 + assert result_generator is not None

299 async for op in result_generator:

300 result = op

301 return TranscriptionResponse(text=result.outputs[0].text)

vllm/v1/engine/async_llm.py

+6

@@ -361,6 +361,12 @@ class AsyncLLM(EngineClient):

361 async def reset_prefix_cache(self) -> None:

362 await self.engine_core.reset_prefix_cache_async()

363

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

365 + await self.engine_core.sleep_async(level)

366 +

367 + async def wake_up(self) -> None:

368 + await self.engine_core.wake_up_async()

369 +

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

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

372 await self.engine_core.add_lora_async(lora_request)

vllm/v1/engine/core.py

+6

@@ -213,6 +213,12 @@ class EngineCore:

213 def reset_prefix_cache(self):

214 self.scheduler.reset_prefix_cache()

215

216 + def sleep(self, level: int = 1):

217 + self.model_executor.sleep(level)

218 +

219 + def wake_up(self):

220 + self.model_executor.wake_up()

221 +

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

223 self.model_executor.add_lora(lora_request)

224

vllm/v1/engine/core_client.py

+30

@@ -81,6 +81,12 @@ class EngineCoreClient(ABC):

81 def reset_prefix_cache(self) -> None:

82 raise NotImplementedError

83

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

85 + raise NotImplementedError

86 +

87 + def wake_up(self) -> None:

88 + raise NotImplementedError

89 +

90 def abort_requests(self, request_ids: List[str]) -> None:

91 raise NotImplementedError

92

@@ -99,6 +105,12 @@ class EngineCoreClient(ABC):

105 async def reset_prefix_cache_async(self) -> None:

106 raise NotImplementedError

107

108 + async def sleep_async(self, level: int = 1) -> None:

109 + raise NotImplementedError

110 +

111 + async def wake_up_async(self) -> None:

112 + raise NotImplementedError

113 +

114 async def abort_requests_async(self, request_ids: List[str]) -> None:

115 raise NotImplementedError

116

@@ -138,6 +150,12 @@ class InprocClient(EngineCoreClient):

150 def reset_prefix_cache(self) -> None:

151 self.engine_core.reset_prefix_cache()

152

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

154 + self.engine_core.sleep(level)

155 +

156 + def wake_up(self) -> None:

157 + self.engine_core.wake_up()

158 +

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

160 self.engine_core.add_lora(lora_request)

161

@@ -307,6 +325,12 @@ class SyncMPClient(MPClient):

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

326 self._call_utility("add_lora", lora_request)

327

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

329 + self._call_utility("sleep", level)

330 +

331 + def wake_up(self) -> None:

332 + self._call_utility("wake_up")

333 +

334

335 class AsyncMPClient(MPClient):

336 """Asyncio-compatible client for multi-proc EngineCore."""

@@ -384,5 +408,11 @@ class AsyncMPClient(MPClient):

408 async def reset_prefix_cache_async(self) -> None:

409 await self._call_utility_async("reset_prefix_cache")

410

411 + async def sleep_async(self, level: int = 1) -> None:

412 + await self._call_utility_async("sleep", level)

413 +

414 + async def wake_up_async(self) -> None:

415 + await self._call_utility_async("wake_up")

416 +

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

418 await self._call_utility_async("add_lora", lora_request)

vllm/v1/engine/llm_engine.py

+6

@@ -169,6 +169,12 @@ class LLMEngine:

169 def reset_prefix_cache(self):

170 self.engine_core.reset_prefix_cache()

171

172 + def sleep(self, level: int = 1):

173 + self.engine_core.sleep(level)

174 +

175 + def wake_up(self):

176 + self.engine_core.wake_up()

177 +

178 def get_tokenizer_group(

179 self,

180 group_type: Type[_G] = BaseTokenizerGroup,

[core] add sleep and wake up endpoint and v1 support (#12987)