| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import traceback |
| 5 | import unittest |
| 6 | |
| 7 | import numpy as np |
| 8 | |
| 9 | from vllm.distributed.device_communicators.shm_object_storage import ( |
| 10 | SingleWriterShmRingBuffer, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | class TestSingleWriterShmRingBuffer(unittest.TestCase): |
| 15 | """Test suite for the ring buffer implementation.""" |
| 16 | |
| 17 | def setUp(self): |
| 18 | """Set up test fixtures.""" |
| 19 | self.buffer_size = 4096 |
| 20 | self.ring_buffer = None |
| 21 | |
| 22 | def tearDown(self): |
| 23 | """Clean up after tests.""" |
| 24 | if self.ring_buffer: |
| 25 | self.ring_buffer.close() |
| 26 | |
| 27 | def test_buffer_opening(self): |
| 28 | """Test opening an existing buffer.""" |
| 29 | # First create a buffer |
| 30 | self.ring_buffer = SingleWriterShmRingBuffer( |
| 31 | data_buffer_size=self.buffer_size, create=True |
| 32 | ) |
| 33 | |
| 34 | # Then open it with another instance |
| 35 | reader_buffer = SingleWriterShmRingBuffer(*self.ring_buffer.handle()) |
| 36 | self.assertFalse(reader_buffer.is_writer) |
| 37 | self.assertEqual( |
| 38 | reader_buffer.shared_memory.name, self.ring_buffer.shared_memory.name |
| 39 | ) |
| 40 | |
| 41 | def test_buffer_access(self): |
| 42 | """Test accessing allocated buffers.""" |
| 43 | self.ring_buffer = SingleWriterShmRingBuffer( |
| 44 | data_buffer_size=self.buffer_size, create=True |
| 45 | ) |
| 46 | |
| 47 | size = 100 |
| 48 | address, monotonic_id = self.ring_buffer.allocate_buf(size) |
| 49 | |
| 50 | # Write some test data |
| 51 | test_data = b"Hello, World!" * 7 # 91 bytes |
| 52 | with self.ring_buffer.access_buf(address) as (data_buf, metadata): |
| 53 | data_buf[0 : len(test_data)] = test_data |
| 54 | |
| 55 | # Read it back |
| 56 | with self.ring_buffer.access_buf(address) as (data_buf2, metadata2): |
| 57 | read_data = bytes(data_buf2[0 : len(test_data)]) |
| 58 | read_id = metadata2[0] |
| 59 | |
| 60 | self.assertEqual(read_data, test_data) |
| 61 | self.assertEqual(read_id, monotonic_id) |
| 62 | |
| 63 | def test_memory_error_on_full_buffer(self): |
| 64 | """Test that MemoryError is raised when buffer is full.""" |
| 65 | small_buffer_size = 200 |
| 66 | self.ring_buffer = SingleWriterShmRingBuffer( |
| 67 | data_buffer_size=small_buffer_size, create=True |
| 68 | ) |
| 69 | |
| 70 | # Fill up the buffer |
| 71 | self.ring_buffer.allocate_buf(100) |
| 72 | self.ring_buffer.allocate_buf(80) # Total: 196 bytes used |
| 73 | |
| 74 | # This should fail |
| 75 | with self.assertRaises(MemoryError): |
| 76 | self.ring_buffer.allocate_buf(1) # Would exceed buffer capacity |
| 77 | |
| 78 | def test_allocation_and_free(self): |
| 79 | """Test allocation and freeing of buffers.""" |
| 80 | small_buffer_size = 200 |
| 81 | self.ring_buffer = SingleWriterShmRingBuffer( |
| 82 | data_buffer_size=small_buffer_size, create=True |
| 83 | ) |
| 84 | |
| 85 | size = 80 |
| 86 | # Write some data |
| 87 | test_data = b"Repeated test data" |
| 88 | for i in range(5): |
| 89 | address, monotonic_id = self.ring_buffer.allocate_buf(size) |
| 90 | with self.ring_buffer.access_buf(address) as (data_buf, metadata): |
| 91 | data_buf[0:4] = (0).to_bytes(4, "little") # 0 for not in-use |
| 92 | data_buf[4 : len(test_data) + 4] = test_data |
| 93 | print(self.ring_buffer.metadata) |
| 94 | freed_ids = self.ring_buffer.free_buf(lambda *args: True) |
| 95 | print(f" Freed IDs: {freed_ids}") |
| 96 | self.assertEqual(freed_ids[0], i) |
| 97 | |
| 98 | def test_clear_buffer(self): |
| 99 | """Test clearing the buffer.""" |
| 100 | self.ring_buffer = SingleWriterShmRingBuffer( |
| 101 | data_buffer_size=self.buffer_size, create=True |
| 102 | ) |
| 103 | |
| 104 | # Allocate some buffers |
| 105 | for _ in range(3): |
| 106 | self.ring_buffer.allocate_buf(100) |
| 107 | |
| 108 | # Clear the buffer |
| 109 | self.ring_buffer.clear() |
| 110 | |
| 111 | # Check that metadata is empty and IDs keep increasing |
| 112 | self.assertEqual(len(self.ring_buffer.metadata), 0) |
| 113 | self.assertEqual(self.ring_buffer.monotonic_id_start, 3) |
| 114 | self.assertEqual(self.ring_buffer.monotonic_id_end, 3) |
| 115 | self.assertEqual(self.ring_buffer.data_buffer_start, 0) |
| 116 | self.assertEqual(self.ring_buffer.data_buffer_end, 0) |
| 117 | |
| 118 | def test_allocation_cycles(self): |
| 119 | buffer_size = 100 |
| 120 | ring = SingleWriterShmRingBuffer(data_buffer_size=buffer_size, create=True) |
| 121 | |
| 122 | # tracking allocations for assertions |
| 123 | allocated_bitmap = np.zeros( |
| 124 | (buffer_size,), dtype=np.bool_ |
| 125 | ) # addr -> is_allocated |
| 126 | allocation_map = dict() # monotonic_id -> (addr, size) |
| 127 | |
| 128 | def count_allocated(bitmap) -> int: |
| 129 | return np.sum(bitmap).item() |
| 130 | |
| 131 | def is_free_fn(a, b) -> bool: |
| 132 | return True |
| 133 | |
| 134 | def mark_allocated_with_assertion(id, addr, size): |
| 135 | addr = addr % buffer_size |
| 136 | self.assertEqual(count_allocated(allocated_bitmap[addr : addr + size]), 0) |
| 137 | |
| 138 | allocated_bitmap[addr : addr + size] = True |
| 139 | allocation_map[id] = (addr, size) |
| 140 | |
| 141 | def mark_freed_with_assertion(id): |
| 142 | self.assertTrue(id in allocation_map) |
| 143 | |
| 144 | addr, size = allocation_map.pop(id) |
| 145 | addr = addr % buffer_size |
| 146 | self.assertEqual( |
| 147 | count_allocated(allocated_bitmap[addr : addr + size]), size |
| 148 | ) |
| 149 | |
| 150 | allocated_bitmap[addr : addr + size] = False |
| 151 | |
| 152 | def ring_free(free_size=None): |
| 153 | freed_ids = ring.free_buf(is_free_fn, free_size) |
| 154 | for freed_id in freed_ids: |
| 155 | mark_freed_with_assertion(freed_id) |
| 156 | |
| 157 | def ring_allocate(allocate_size): |
| 158 | allocate_size_with_md = allocate_size + ring.MD_SIZE |
| 159 | try: |
| 160 | addr, monotonic_id = ring.allocate_buf(allocate_size) |
| 161 | mark_allocated_with_assertion(monotonic_id, addr, allocate_size_with_md) |
| 162 | except MemoryError: |
| 163 | # free 2x size for enough space if wrapping happened |
| 164 | ring_free(allocate_size_with_md * 2) |
| 165 | |
| 166 | # retry allocating |
| 167 | addr, monotonic_id = ring.allocate_buf(allocate_size) |
| 168 | mark_allocated_with_assertion(monotonic_id, addr, allocate_size_with_md) |
| 169 | |
| 170 | # 1. allocation & free cycles |
| 171 | for _ in range(33): |
| 172 | # will consume 2 + 8 = 10 bytes per allocation |
| 173 | ring_allocate(2) |
| 174 | |
| 175 | # 2. free all allocations |
| 176 | ring_free() |
| 177 | |
| 178 | # 3. try allocate the largest possible buffer |
| 179 | ring_allocate(buffer_size - ring.MD_SIZE) |
| 180 | |
| 181 | |
| 182 | def main(): |
| 183 | """Main function demonstrating usage and running tests.""" |
| 184 | print("=== SingleWriterShmRingBuffer Test Suite ===\n") |
| 185 | |
| 186 | # Run unit tests |
| 187 | print("Running unit tests...") |
| 188 | unittest.main(argv=[""], exit=False, verbosity=2) |
| 189 | |
| 190 | print("\n" + "=" * 50) |
| 191 | print("=== Manual Demo ===\n") |
| 192 | |
| 193 | # Manual demonstration |
| 194 | try: |
| 195 | print("Creating ring buffer...") |
| 196 | writer_buffer = SingleWriterShmRingBuffer(data_buffer_size=2048, create=True) |
| 197 | reader_buffer = SingleWriterShmRingBuffer(*writer_buffer.handle()) |
| 198 | |
| 199 | print(f"Buffer created with name: {writer_buffer.shared_memory.name}") |
| 200 | |
| 201 | # Allocate some buffers |
| 202 | print("\nAllocating buffers...") |
| 203 | address_array = [] |
| 204 | for i in range(3): |
| 205 | size = 100 + i * 50 |
| 206 | try: |
| 207 | writer_buffer.free_buf(lambda *args: True) |
| 208 | address, monotonic_id = writer_buffer.allocate_buf(size) |
| 209 | address_array.append((address, size, monotonic_id)) |
| 210 | |
| 211 | # Write some test data |
| 212 | with writer_buffer.access_buf(address) as (data_buf, metadata): |
| 213 | test_message = f"Test message {i}".encode() |
| 214 | data_buf[0 : len(test_message)] = test_message |
| 215 | |
| 216 | except MemoryError as e: |
| 217 | print(f" Failed to allocate {size} bytes: {e}") |
| 218 | |
| 219 | print("\nBuffer state:") |
| 220 | print(f" Data buffer start: {writer_buffer.data_buffer_start}") |
| 221 | print(f" Data buffer end: {writer_buffer.data_buffer_end}") |
| 222 | print(f" Monotonic ID start: {writer_buffer.monotonic_id_start}") |
| 223 | print(f" Monotonic ID end: {writer_buffer.monotonic_id_end}") |
| 224 | print(f" Metadata entries: {len(writer_buffer.metadata)}") |
| 225 | |
| 226 | # Try to read back the data |
| 227 | print("\nReading back data...") |
| 228 | for address, size, monotonic_id in address_array: |
| 229 | with reader_buffer.access_buf(address) as (data_buf, metadata): |
| 230 | # Find null terminator or read first 50 chars |
| 231 | data_bytes = bytes(data_buf[0:size]) |
| 232 | message = data_bytes.decode() |
| 233 | print(f" ID {monotonic_id}: '{message}'") |
| 234 | |
| 235 | except Exception as e: |
| 236 | print(f"Demo error: {e}") |
| 237 | traceback.print_exc() |
| 238 | |
| 239 | print("\n=== Demo Complete ===") |
| 240 | |
| 241 | |
| 242 | if __name__ == "__main__": |
| 243 | main() |