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 9, 23:00 UTC
[Core] Cleanup shm based object store on engine shutdown (#32429)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
Walter Beller-Morales committed
Jan 20, 2026 at 03:53 UTC
8be263c3fb1f98d85bd6a06d52e6036057f8814e
7 files changed
+57
-7
tests/distributed/test_shm_buffer.py
+1
-1
@@ -22,7 +22,7 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase):
22
def tearDown(self):
23
"""Clean up after tests"""
24
if self.ring_buffer:
25
- del self.ring_buffer
25
+ self.ring_buffer.close()
26
27
def test_buffer_opening(self):
28
"""Test opening an existing buffer"""
tests/distributed/test_shm_storage.py
+1
-1
@@ -56,7 +56,7 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase):
56
def tearDown(self):
57
"""Clean up after each test."""
58
if self.storage:
59
- del self.storage
59
+ self.storage.close()
60
61
def test_minimal_put_get_cycle(self):
62
"""Test basic put and get operations."""
vllm/distributed/device_communicators/shm_object_storage.py
+13
-3
@@ -4,7 +4,7 @@
4
import pickle
5
from abc import ABC, abstractmethod
6
from collections.abc import Callable, Iterable
7
-from contextlib import contextmanager
7
+from contextlib import contextmanager, suppress
8
from dataclasses import dataclass
9
from itertools import chain
10
from multiprocessing import shared_memory
@@ -126,6 +126,7 @@ class SingleWriterShmRingBuffer:
126
self.data_buffer_end = 0
127
128
if create:
129
+ logger.debug("Creating new shared memory buffer: %s", name)
130
# we are creating a buffer
131
self.metadata: dict[int, int] = {} # monotonic_id -> start address
132
self.shared_memory = shared_memory.SharedMemory(
@@ -169,11 +170,16 @@ class SingleWriterShmRingBuffer:
170
self.data_buffer_start = 0
171
self.data_buffer_end = 0
172
172
- def __del__(self):
173
+ def close(self) -> None:
174
+ """Close the shared memory."""
175
if hasattr(self, "shared_memory"):
176
self.shared_memory.close()
177
if self.is_writer:
176
- self.shared_memory.unlink()
178
+ with suppress(FileNotFoundError):
179
+ self.shared_memory.unlink()
180
+
181
+ def __del__(self):
182
+ self.close()
183
184
def int2byte(self, integer: int) -> bytes:
185
"""Convert an integer to bytes."""
@@ -663,6 +669,10 @@ class SingleWriterShmObjectStorage:
669
if reader_count >= self.n_readers:
670
self.increment_reader_flag(data_view[: self.flag_bytes])
671
672
+ def close(self) -> None:
673
+ """Close the shared memory."""
674
+ self.ring_buffer.close()
675
+
676
def handle(self):
677
"""Get handle for sharing across processes."""
678
return ShmObjectStorageHandle(
vllm/envs.py
+27
-2
@@ -7,6 +7,7 @@ import logging
7
import os
8
import sys
9
import tempfile
10
+import uuid
11
from collections.abc import Callable
12
from typing import TYPE_CHECKING, Any, Literal
13
@@ -457,6 +458,27 @@ def get_vllm_port() -> int | None:
458
raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err
459
460
461
+def get_env_or_set_default(
462
+ env_name: str,
463
+ default_factory: Callable[[], str],
464
+) -> Callable[[], str]:
465
+ """
466
+ Create a lambda that returns an environment variable value if set,
467
+ or generates and sets a default value using the provided factory function.
468
+ """
469
+
470
+ def _get_or_set_default() -> str:
471
+ value = os.getenv(env_name)
472
+ if value is not None:
473
+ return value
474
+
475
+ default_value = default_factory()
476
+ os.environ[env_name] = default_value
477
+ return default_value
478
+
479
+ return _get_or_set_default
480
+
481
+
482
# The start-* and end* here are used by the documentation generator
483
# to extract the used env vars.
484
@@ -1558,8 +1580,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1580
),
1581
# Name of the shared memory buffer used for object storage.
1582
# Only effective when mm_config.mm_processor_cache_type == "shm".
1561
- "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME": lambda: os.getenv(
1562
- "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", "VLLM_OBJECT_STORAGE_SHM_BUFFER"
1583
+ # Automatically generates a unique UUID-based name per process tree
1584
+ # if not explicitly set.
1585
+ "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME": get_env_or_set_default(
1586
+ "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME",
1587
+ lambda: f"VLLM_OBJECT_STORAGE_SHM_BUFFER_{uuid.uuid4().hex}",
1588
),
1589
# The size in MB of the buffers (NVL and RDMA) used by DeepEP
1590
"VLLM_DEEPEP_BUFFER_SIZE_MB": lambda: int(
vllm/multimodal/cache.py
+8
@@ -295,6 +295,10 @@ class BaseMultiModalProcessorCache(
295
"""
296
return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes]
297
298
+ def close(self) -> None:
299
+ """Close the underlying cache, if needed."""
300
+ pass
301
+
302
@abstractmethod
303
def touch_sender_cache_item(self, mm_hash: str) -> None:
304
"""
@@ -534,6 +538,10 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
538
def make_stats(self, *, delta: bool = False) -> CacheInfo:
539
return self._stat(delta=delta)
540
541
+ @override
542
+ def close(self) -> None:
543
+ self._shm_cache.close()
544
+
545
def remove_dangling_items(self) -> None:
546
"""Remove items that are no longer in the shared memory cache."""
547
cached_hashes = self._shm_cache.key_index.keys()
vllm/v1/engine/async_llm.py
+3
@@ -249,6 +249,9 @@ class AsyncLLM(EngineClient):
249
if engine_core := getattr(self, "engine_core", None):
250
engine_core.shutdown()
251
252
+ if input_processor := getattr(self, "input_processor", None):
253
+ input_processor.close()
254
+
255
handler = getattr(self, "output_handler", None)
256
if handler is not None:
257
cancel_task_threadsafe(handler)
vllm/v1/engine/input_processor.py
+4
@@ -712,3 +712,7 @@ class InputProcessor:
712
713
def clear_mm_cache(self) -> None:
714
self.input_preprocessor.clear_mm_cache()
715
+
716
+ def close(self) -> None:
717
+ if self.mm_processor_cache is not None:
718
+ self.mm_processor_cache.close()