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, 00:30 UTC
[Fix] Skip `record_sleep_state` logic in `PrometheusStatsLogger` if not in dev mode (#27789)
Signed-off-by: SumanthRH <sumanthrh99@gmail.com>
Sumanth R Hegde committed
Oct 30, 2025 at 12:26 UTC
4917002523db90813a47ca5aed5cd22e2edb75f4
2 files changed
+45
-1
tests/basic_correctness/test_cumem.py
+42
-1
@@ -1,10 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import asyncio
5
+
6
import pytest
7
import torch
8
7
-from vllm import LLM, SamplingParams
9
+from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
10
from vllm.device_allocator.cumem import CuMemAllocator
11
from vllm.utils.mem_constants import GiB_bytes
12
@@ -201,3 +203,42 @@ def test_deep_sleep():
203
204
# cmp output
205
assert output[0].outputs[0].text == output2[0].outputs[0].text
206
+
207
+
208
+@create_new_process_for_each_test()
209
+def test_deep_sleep_async():
210
+ async def test():
211
+ model = "hmellor/tiny-random-LlamaForCausalLM"
212
+ free, total = torch.cuda.mem_get_info()
213
+ used_bytes_baseline = total - free # in case other process is running
214
+ engine_args = AsyncEngineArgs(
215
+ model=model,
216
+ enable_sleep_mode=True,
217
+ )
218
+
219
+ llm = AsyncLLMEngine.from_engine_args(engine_args)
220
+ prompt = "How are you?"
221
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
222
+ outputs = llm.generate(prompt, sampling_params, request_id="test_request_id1")
223
+ async for output in outputs:
224
+ pass
225
+
226
+ # Put the engine to deep sleep
227
+ await llm.sleep(level=2)
228
+
229
+ await llm.wake_up(tags=["weights"])
230
+ await llm.collective_rpc("reload_weights")
231
+ free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
232
+ used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
233
+ assert used_bytes < 4 * GiB_bytes
234
+
235
+ # now allocate kv cache and cuda graph memory
236
+ await llm.wake_up(tags=["kv_cache"])
237
+ outputs2 = llm.generate(prompt, sampling_params, request_id="test_request_id2")
238
+ async for output2 in outputs2:
239
+ pass
240
+
241
+ # cmp output
242
+ assert output.outputs[0].text == output2.outputs[0].text
243
+
244
+ asyncio.run(test())
vllm/v1/metrics/loggers.py
+3
@@ -1052,6 +1052,9 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
1052
self.gauge_lora_info.labels(**lora_info_labels).set_to_current_time()
1053
1054
def record_sleep_state(self, sleep: int = 0, level: int = 0):
1055
+ if not envs.VLLM_SERVER_DEV_MODE:
1056
+ return
1057
+
1058
awake = 1
1059
discard_all = 0
1060
weights_offloaded = 0