diff --git a/tests/distributed/test_shm_storage.py b/tests/distributed/test_shm_storage.py index 762ffad8e6..66c2751708 100644 --- a/tests/distributed/test_shm_storage.py +++ b/tests/distributed/test_shm_storage.py @@ -291,6 +291,93 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase): reader.touch("key", address, monotonic_id, signature) +class TestShmObjectStorageReferenceCounting(unittest.TestCase): + """One writer and four readers, as with tensor-parallel workers.""" + + def setUp(self): + ring_buffer = SingleWriterShmRingBuffer( + data_buffer_size=1024 * 100, create=True + ) + self.writer = SingleWriterShmObjectStorage( + max_object_size=1024 * 10, + n_readers=4, + ring_buffer=ring_buffer, + serde_class=MsgpackSerde, + ) + self.addCleanup(self.writer.close) + handle = self.writer.handle() + handle.reader_lock = Lock() + self.readers = [ + SingleWriterShmObjectStorage.create_from_handle(handle) for _ in range(4) + ] + for reader in self.readers: + self.addCleanup(reader.close) + + def _reference(self, key: str) -> tuple[int, int, list[int], str]: + """Take one writer reference to the key, as one request item does.""" + if self.writer.is_cached(key): + address, monotonic_id = self.writer.get_cached(key) + else: + address, monotonic_id = self.writer.put(key, "value") + return address, monotonic_id, self.writer.get_signature(key), key + + @staticmethod + def _touch(reader: SingleWriterShmObjectStorage, reference) -> None: + address, monotonic_id, signature, key = reference + reader.touch(key, address, monotonic_id, signature) + + def _is_evictable(self, key: str) -> bool: + self.writer.free_unused() + return not self.writer.is_cached(key) + + def test_item_referenced_again_before_every_reader_fetched_it(self): + """A second request must not leave the item unevictable forever.""" + first = self._reference("image") + self.writer.touch("image") + second = self._reference("image") + self.writer.release_touches() + + # Every reader starts on the first request; the ones that finish it + # go on to the second while the others are still fetching. + for reader in self.readers: + self._touch(reader, first) + for reader in self.readers: + reader.get(*first) + self._touch(reader, second) + reader.get(*second) + + self.assertTrue(self._is_evictable("image")) + + def test_new_item_twice_in_one_request_is_not_evicted_early(self): + """The item must stay until the last reader has fetched it.""" + self.writer.touch("image") + references = [self._reference("image"), self._reference("image")] + self.writer.release_touches() + + for reader in self.readers[:-1]: + for reference in references: + self._touch(reader, reference) + for reference in references: + reader.get(*reference) + for reference in references: + self._touch(self.readers[-1], reference) + self.assertFalse(self._is_evictable("image")) + + for reference in references: + self.readers[-1].get(*reference) + self.assertTrue(self._is_evictable("image")) + + def test_touch_protects_item_until_released(self): + reference = self._reference("image") + for reader in self.readers: + reader.get(*reference) + + self.writer.touch("image") + self.assertFalse(self._is_evictable("image")) + self.writer.release_touches() + self.assertTrue(self._is_evictable("image")) + + # Reader process function def reader_process(process_id, storage_handle, items_to_read): """Reader process that connects to existing shared memory and reads data.""" diff --git a/vllm/distributed/device_communicators/shm_object_storage.py b/vllm/distributed/device_communicators/shm_object_storage.py index 3da8c5dcf6..081fbea197 100644 --- a/vllm/distributed/device_communicators/shm_object_storage.py +++ b/vllm/distributed/device_communicators/shm_object_storage.py @@ -527,6 +527,9 @@ class SingleWriterShmObjectStorage: self.id_index: dict[int, str] = {} # Writer flag to track in-use status: monotonic_id -> count self.writer_flag: dict[int, int] = {} + # Items touch() keeps from being freed until they are looked up + # or release_touches() is called + self._touched: set[int] = set() else: if reader_lock is None: raise ValueError("Lock must be provided for readers.") @@ -540,6 +543,7 @@ class SingleWriterShmObjectStorage: self.key_index.clear() self.id_index.clear() self.writer_flag.clear() + self._touched.clear() logger.debug("Object storage cleared and reinitialized.") def copy_to_buffer( @@ -664,8 +668,13 @@ class SingleWriterShmObjectStorage: """Get the cached object by key if it exists.""" address, monotonic_id = self.key_index[key] self.increment_writer_flag(monotonic_id) + self._touched.discard(monotonic_id) return address, monotonic_id + def release_touches(self) -> None: + """Stop protecting touched items that were not looked up.""" + self._touched.clear() + def get_signature(self, key: str) -> list[int]: """Sign the handle of a cached object so readers can verify it.""" address, monotonic_id = self.key_index[key] @@ -756,10 +765,11 @@ class SingleWriterShmObjectStorage: monotonic_id: int = 0, signature: list[int] | None = None, ) -> None: - """Touch an existing cached item to update its eviction status. + """Touch an existing cached item. - For writers (ShmObjectStoreSenderCache): Increment writer_flag - For readers (ShmObjectStoreReceiverCache): Increment reader_count + For writers (ShmObjectStoreSenderCache): Protect the item from + eviction until get_cached or release_touches + For readers (ShmObjectStoreReceiverCache): Validate the handle Args: key: String key of the object to touch @@ -771,25 +781,12 @@ class SingleWriterShmObjectStorage: if self._reader_lock is None: if key not in self.key_index: return None - address, monotonic_id = self.key_index[key] - # Writer side: increment writer_flag to raise eviction threshold - self.increment_writer_flag(monotonic_id) + self._touched.add(self.key_index[key][1]) else: # Reject forged handles before dereferencing the supplied address. self.verify_signature(key, address, monotonic_id, signature) - with ( - self._reader_lock, - self.ring_buffer.access_buf(address) as (data_view, buf_metadata), - ): + with self.ring_buffer.access_buf(address) as (_, buf_metadata): self._validate_monotonic_id(address, monotonic_id, buf_metadata) - reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes]) - - # NOTE(Long): - # Avoid increasing flag on newly added item (sync with sender) - # Since when a new item is added - # pre-touch has no effect on writer side - if reader_count >= self.n_readers: - self.increment_reader_flag(data_view[: self.flag_bytes]) def close(self) -> None: """Close the shared memory.""" @@ -823,6 +820,8 @@ class SingleWriterShmObjectStorage: """Default is_free function that checks if the first 4 bytes are zero. This indicates that the buffer is free. """ + if id in self._touched: + return False reader_count = int.from_bytes(buf[0:4], "little", signed=True) writer_count = self.writer_flag[id] return reader_count >= writer_count * self.n_readers diff --git a/vllm/multimodal/cache/base.py b/vllm/multimodal/cache/base.py index d2aaeb50e3..f613325f85 100644 --- a/vllm/multimodal/cache/base.py +++ b/vllm/multimodal/cache/base.py @@ -377,6 +377,10 @@ class BaseMultiModalProcessorCache( """ raise NotImplementedError + def release_sender_touches(self) -> None: + """Release the items that were touched but not updated afterwards.""" + return None + @abstractmethod def make_stats(self, *, delta: bool = False) -> CacheInfo: """Get (and reset) the multi-modal cache stats. diff --git a/vllm/multimodal/cache/shm.py b/vllm/multimodal/cache/shm.py index ab84dcf6af..2fd8ad47e3 100644 --- a/vllm/multimodal/cache/shm.py +++ b/vllm/multimodal/cache/shm.py @@ -170,10 +170,13 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): @override def touch_sender_cache_item(self, mm_hash: str) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments writer_flag on sender side.""" + """Touch the item in shared memory cache to prevent eviction.""" self._shm_cache.touch(mm_hash) + @override + def release_sender_touches(self) -> None: + self._shm_cache.release_touches() + @override def validate_input_item( self, @@ -305,8 +308,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): mm_hash: str, mm_item: MultiModalKwargsItem | None = None, ) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments reader_count on receiver side.""" + """Validate the item's handle in shared memory cache.""" assert mm_item is not None if (handle := _get_shm_handle(mm_item)) is not None: address, monotonic_id, signature = handle diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index eb53023f0f..7267020353 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1534,13 +1534,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]): ) with timing_ctx.record("merge_mm_kwargs"): - mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( - cache, - mm_hashes=mm_hashes, - mm_is_cached=mm_is_cached, - mm_missing_kwargs=mm_missing_kwargs, - mm_missing_prompt_updates=mm_missing_prompt_updates, - ) + try: + mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( + cache, + mm_hashes=mm_hashes, + mm_is_cached=mm_is_cached, + mm_missing_kwargs=mm_missing_kwargs, + mm_missing_prompt_updates=mm_missing_prompt_updates, + ) + finally: + cache.release_sender_touches() return MultiModalProcessingResult( prompt_ids=self._postprocess_prompt(inputs.prompt),