| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import hashlib |
| 5 | import hmac |
| 6 | import pickle |
| 7 | import secrets |
| 8 | from abc import ABC, abstractmethod |
| 9 | from collections.abc import Callable, Iterable |
| 10 | from contextlib import contextmanager, suppress |
| 11 | from dataclasses import dataclass |
| 12 | from itertools import chain |
| 13 | from multiprocessing import shared_memory |
| 14 | from multiprocessing.synchronize import Lock as LockType |
| 15 | from typing import Any |
| 16 | from unittest.mock import patch |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.logger import init_logger |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | class SingleWriterShmRingBuffer: |
| 26 | """A single-writer, multiple-reader ring buffer implementation using shared |
| 27 | memory. This class provides a thread-safe ring buffer where one process |
| 28 | can write data while multiple processes/threads can read from it. |
| 29 | |
| 30 | Architecture: |
| 31 | - Uses shared memory for cross-process communication |
| 32 | - Maintains metadata for each allocated buffer chunk in the writer process |
| 33 | - Supports custom "is_free_fn" functions to determine when buffers can be |
| 34 | reused |
| 35 | - Each buffer chunk contains: `[4-byte id][4-byte size][actual_data]` |
| 36 | |
| 37 | Key Concepts: |
| 38 | - monotonic_id_start/end: Track the range of active buffer IDs |
| 39 | - data_buffer_start/end: Track the physical memory range in use |
| 40 | - Automatic wraparound when reaching buffer end |
| 41 | - Lazy garbage collection based on is_free_fn checks |
| 42 | |
| 43 | Example Usage Scenarios: |
| 44 | |
| 45 | Scenario 1: Simple Linear Allocation |
| 46 | ``` |
| 47 | Buffer size: 100 bytes |
| 48 | Initial state: [................................................. ] |
| 49 | ^start=end(0) |
| 50 | |
| 51 | After allocating 20 bytes (id=0): |
| 52 | [id:0|size:20|data........][...................................] |
| 53 | ^start(0) ^end(28) |
| 54 | |
| 55 | After allocating 30 bytes (id=1): |
| 56 | [id:0|size:20|data........][id:1|size:30|data..............][..] |
| 57 | ^start(0) ^end(66) |
| 58 | ``` |
| 59 | |
| 60 | Scenario 2: Memory Reclamation |
| 61 | ``` |
| 62 | Before freeing (both buffers still in use): |
| 63 | [id:0|size:20|data........][id:1|size:30|data..............][..] |
| 64 | ^start(0) ^end(66) |
| 65 | |
| 66 | After id:0 is marked free by readers: |
| 67 | [FREED.................... ][id:1|size:30|data..............][..] |
| 68 | ^start(28) ^end(66) |
| 69 | |
| 70 | After both are freed: |
| 71 | [FREED..............................................][..] |
| 72 | ^start=end(66) |
| 73 | ``` |
| 74 | |
| 75 | Scenario 3: Wraparound Allocation (continuing from Scenario 2) |
| 76 | ``` |
| 77 | Starting from after memory reclamation in Scenario 2: |
| 78 | [FREED..............................................][..] |
| 79 | ^start=end(66) |
| 80 | |
| 81 | Allocate 40 bytes (id=2) - only 34 bytes available at end, so wraparound: |
| 82 | [id:2|size:40|data........................][FREED.............][..] |
| 83 | ^end(148) ^start(66) |
| 84 | ``` |
| 85 | |
| 86 | Scenario 4: Error Handling - Out of Space |
| 87 | ``` |
| 88 | Starting from after wraparound allocation in Scenario 3: |
| 89 | [id:2|size:40|data........................][FREED.............][..] |
| 90 | ^end(148) ^start(66) |
| 91 | |
| 92 | Trying to allocate 20 more bytes: |
| 93 | occupied_size_new = end + size - start = 148 + 28 - 66 > buffer_size(100) |
| 94 | -> Raises MemoryError: "Not enough space in the data buffer" |
| 95 | ``` |
| 96 | |
| 97 | Thread Safety: |
| 98 | - Single writer: Only one process/thread should write (allocate_buf) |
| 99 | - Multiple readers: Multiple processes/threads can read (access_buf) |
| 100 | - Reader synchronization handled by is_free_fn callback |
| 101 | - Writer handles garbage collection (free_buf) based on reader feedback |
| 102 | |
| 103 | Shared Memory Layout: |
| 104 | `[32-byte handle-signing secret][ring buffer chunks...]` |
| 105 | |
| 106 | Memory Layout per Buffer Chunk: |
| 107 | `[4-byte monotonic_id][4-byte chunk_size][actual_data...]` |
| 108 | ^metadata_start ^data_start |
| 109 | |
| 110 | The monotonic_id ensures data integrity - readers can verify they're |
| 111 | accessing the correct data even after buffer wraparound or reuse. |
| 112 | """ |
| 113 | |
| 114 | def __init__( |
| 115 | self, |
| 116 | data_buffer_size: int, |
| 117 | name: str | None = None, |
| 118 | create: bool = False, |
| 119 | ): |
| 120 | self.data_buffer_size = data_buffer_size |
| 121 | self.is_writer = create |
| 122 | |
| 123 | self.ID_NBYTES = 4 |
| 124 | self.ID_MAX = 2**31 # exclusive, so 2**31 - 1 is the max value |
| 125 | self.SIZE_NBYTES = 4 |
| 126 | # 4 bytes for id, 4 bytes for buffer size |
| 127 | self.MD_SIZE = self.ID_NBYTES + self.SIZE_NBYTES |
| 128 | self.AUTH_SECRET_BYTES = 32 |
| 129 | self.DATA_OFFSET = self.AUTH_SECRET_BYTES |
| 130 | self.monotonic_id_end = 0 |
| 131 | self.monotonic_id_start = 0 |
| 132 | self.data_buffer_start = 0 |
| 133 | self.data_buffer_end = 0 |
| 134 | |
| 135 | if create: |
| 136 | logger.debug("Creating new shared memory buffer: %s", name) |
| 137 | # we are creating a buffer |
| 138 | self.metadata: dict[int, int] = {} # monotonic_id -> start address |
| 139 | self.shared_memory = shared_memory.SharedMemory( |
| 140 | create=True, |
| 141 | size=self.data_buffer_size + self.DATA_OFFSET, |
| 142 | name=name, |
| 143 | ) |
| 144 | self._init_auth_secret() |
| 145 | else: |
| 146 | # we are opening an existing buffer |
| 147 | # fix to https://stackoverflow.com/q/62748654/9191338 |
| 148 | # Python incorrectly tracks shared memory even if it is not |
| 149 | # created by the process. The following patch is a workaround. |
| 150 | with patch( |
| 151 | "multiprocessing.resource_tracker.register", |
| 152 | lambda *args, **kwargs: None, |
| 153 | ): |
| 154 | self.shared_memory = shared_memory.SharedMemory(name=name) |
| 155 | # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa |
| 156 | # Some platforms allocate memory based on page size, |
| 157 | # so the shared memory block size may be larger or equal |
| 158 | # to the requested size. The size parameter is ignored |
| 159 | # when attaching to an existing block. |
| 160 | assert ( |
| 161 | self.shared_memory.size >= self.data_buffer_size + self.DATA_OFFSET |
| 162 | ) |
| 163 | |
| 164 | logger.debug( |
| 165 | "Shared memory created/opened with name: %s, size: %d", |
| 166 | self.shared_memory.name, |
| 167 | self.data_buffer_size, |
| 168 | ) |
| 169 | |
| 170 | def handle(self): |
| 171 | return ( |
| 172 | self.data_buffer_size, |
| 173 | self.shared_memory.name, |
| 174 | ) |
| 175 | |
| 176 | def clear(self) -> None: |
| 177 | """Clear the ring buffer.""" |
| 178 | assert self.is_writer, "Only the writer can clear the buffer." |
| 179 | self.metadata.clear() |
| 180 | # Keep IDs increasing so handles issued before the clear (e.g. for |
| 181 | # requests still draining) fail the ID check once their chunk is |
| 182 | # reused, instead of resolving to a different object. |
| 183 | self.monotonic_id_start = self.monotonic_id_end |
| 184 | self.data_buffer_start = 0 |
| 185 | self.data_buffer_end = 0 |
| 186 | |
| 187 | def get_auth_secret(self) -> bytes: |
| 188 | """Return the current handle-signing secret.""" |
| 189 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 190 | return bytes(self.shared_memory.buf[: self.AUTH_SECRET_BYTES]) |
| 191 | |
| 192 | def _init_auth_secret(self) -> None: |
| 193 | assert self.is_writer, "Only the writer can set the auth secret." |
| 194 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 195 | self.shared_memory.buf[: self.AUTH_SECRET_BYTES] = secrets.token_bytes( |
| 196 | self.AUTH_SECRET_BYTES |
| 197 | ) |
| 198 | |
| 199 | def close(self) -> None: |
| 200 | """Close the shared memory.""" |
| 201 | if hasattr(self, "shared_memory"): |
| 202 | self.shared_memory.close() |
| 203 | if self.is_writer: |
| 204 | with suppress(FileNotFoundError): |
| 205 | self.shared_memory.unlink() |
| 206 | |
| 207 | def __del__(self): |
| 208 | self.close() |
| 209 | |
| 210 | def int2byte(self, integer: int) -> bytes: |
| 211 | """Convert an integer to bytes.""" |
| 212 | return integer.to_bytes(self.ID_NBYTES, "little", signed=True) |
| 213 | |
| 214 | def byte2int(self, byte_data: bytes) -> int: |
| 215 | """Convert bytes back to an integer.""" |
| 216 | return int.from_bytes(byte_data, "little", signed=True) |
| 217 | |
| 218 | def allocate_buf(self, size: int) -> tuple[int, int]: |
| 219 | """Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory. |
| 220 | Memory layout: |
| 221 | `[4-byte monotonic_id][4-byte size][buffer data...]` |
| 222 | """ |
| 223 | assert self.is_writer, "Only the writer can allocate buffers." |
| 224 | assert size > 0, "Size must be greater than 0" |
| 225 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 226 | size += self.MD_SIZE # add metadata size to the buffer size |
| 227 | # reset to beginning if the buffer does have enough contiguous space |
| 228 | buffer_end_reset = self.data_buffer_end % self.data_buffer_size |
| 229 | if buffer_end_reset + size > self.data_buffer_size: |
| 230 | buffer_end_reset = ( |
| 231 | self.data_buffer_end // self.data_buffer_size + 1 |
| 232 | ) * self.data_buffer_size |
| 233 | else: # no reset needed |
| 234 | buffer_end_reset = self.data_buffer_end |
| 235 | |
| 236 | # check if we have enough space in the data buffer |
| 237 | # i.e. if the new end (self.data_buffer_end + size) |
| 238 | # exceeds the start of the data buffer |
| 239 | occupied_size_new = buffer_end_reset + size - self.data_buffer_start |
| 240 | if occupied_size_new > self.data_buffer_size: |
| 241 | raise MemoryError( |
| 242 | "Not enough space in the data buffer, " |
| 243 | "try calling free_buf() to free up space" |
| 244 | ) |
| 245 | self.data_buffer_end = buffer_end_reset |
| 246 | |
| 247 | # first 4 bytes as the monotonic id |
| 248 | buf_idx = self.data_buffer_end % self.data_buffer_size |
| 249 | physical_idx = self.DATA_OFFSET + buf_idx |
| 250 | self.shared_memory.buf[physical_idx : physical_idx + self.ID_NBYTES] = ( |
| 251 | self.int2byte(self.monotonic_id_end) |
| 252 | ) |
| 253 | # next 4 bytes as the size of the data buffer |
| 254 | self.shared_memory.buf[ |
| 255 | physical_idx + self.ID_NBYTES : physical_idx + self.MD_SIZE |
| 256 | ] = self.int2byte(size) |
| 257 | |
| 258 | # record metadata |
| 259 | self.metadata[self.monotonic_id_end % self.ID_MAX] = self.data_buffer_end |
| 260 | # update buffer and monotonic id indices |
| 261 | current_buffer_end = self.data_buffer_end |
| 262 | current_id_end = self.monotonic_id_end |
| 263 | self.data_buffer_end += size |
| 264 | self.monotonic_id_end = (self.monotonic_id_end + 1) % self.ID_MAX |
| 265 | return current_buffer_end, current_id_end |
| 266 | |
| 267 | @contextmanager |
| 268 | def access_buf(self, address: int): |
| 269 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 270 | buf_idx = address % self.data_buffer_size |
| 271 | physical_idx = self.DATA_OFFSET + buf_idx |
| 272 | |
| 273 | # read metadata |
| 274 | metadata_buff = self.shared_memory.buf[ |
| 275 | physical_idx : physical_idx + self.MD_SIZE |
| 276 | ] |
| 277 | id = self.byte2int(metadata_buff[: self.ID_NBYTES]) |
| 278 | size = self.byte2int(metadata_buff[self.ID_NBYTES : self.MD_SIZE]) |
| 279 | |
| 280 | # yield the data buffer and metadata |
| 281 | data_buff = self.shared_memory.buf[ |
| 282 | physical_idx + self.MD_SIZE : physical_idx + size |
| 283 | ] |
| 284 | with ( |
| 285 | memoryview(data_buff) as data_view, |
| 286 | ): |
| 287 | yield data_view, (id, size) |
| 288 | |
| 289 | def free_buf( |
| 290 | self, |
| 291 | is_free_fn: Callable[[int, memoryview], bool], |
| 292 | nbytes: int | None = None, |
| 293 | ) -> Iterable[int]: |
| 294 | """Free a buffer of the given size. This is a no-op in shared memory, |
| 295 | but we need to keep track of the metadata. |
| 296 | |
| 297 | If freed memory spreads across the end and start of the ring buffer, |
| 298 | the actual freed memory will be in two segments. In this case there |
| 299 | still might not be a contiguous space of `nbytes` available. |
| 300 | |
| 301 | Args: |
| 302 | is_free_fn (Callable[[int, memoryview], bool]): Predicate called with |
| 303 | a monotonic id and the buffer, returning True when that buffer |
| 304 | can be reclaimed. |
| 305 | nbytes (int, optional): The size of the buffer to free. If None, |
| 306 | frees the maximum size of the ring buffer. |
| 307 | |
| 308 | """ |
| 309 | assert self.is_writer, "Only the writer can free buffers." |
| 310 | logger.debug( |
| 311 | "Freeing up space in the ring buffer, " |
| 312 | "monotonic_id_start: %d, monotonic_id_end: %d", |
| 313 | self.monotonic_id_start, |
| 314 | self.monotonic_id_end, |
| 315 | ) |
| 316 | monotonic_id_before = self.monotonic_id_start |
| 317 | # if nbytes is None, free up the maximum size of the ring buffer |
| 318 | if nbytes is None: |
| 319 | nbytes = self.data_buffer_size |
| 320 | freed_bytes = 0 |
| 321 | while self.monotonic_id_start in self.metadata and freed_bytes < nbytes: |
| 322 | address = self.metadata[self.monotonic_id_start] |
| 323 | with self.access_buf(address) as (data_buff, metadata): |
| 324 | if is_free_fn(self.monotonic_id_start, data_buff): |
| 325 | # check passed, we can free the buffer |
| 326 | del self.metadata[self.monotonic_id_start] |
| 327 | self.monotonic_id_start = ( |
| 328 | self.monotonic_id_start + 1 |
| 329 | ) % self.ID_MAX |
| 330 | if self.monotonic_id_start in self.metadata: |
| 331 | # pointing to the start addr of next allocation |
| 332 | self.data_buffer_start += ( |
| 333 | self.metadata[self.monotonic_id_start] |
| 334 | - self.data_buffer_start |
| 335 | ) % self.data_buffer_size |
| 336 | else: |
| 337 | # no remaining allocation, reset to zero |
| 338 | self.data_buffer_start = self.data_buffer_end = 0 |
| 339 | freed_bytes += metadata[1] |
| 340 | else: |
| 341 | # there are still readers, we cannot free the buffer |
| 342 | break |
| 343 | |
| 344 | logger.debug( |
| 345 | "Freed %d bytes from the ring buffer, " |
| 346 | "monotonic_id_start: %d, monotonic_id_end: %d", |
| 347 | freed_bytes, |
| 348 | self.monotonic_id_start, |
| 349 | self.monotonic_id_end, |
| 350 | ) |
| 351 | |
| 352 | # buffer wrap around |
| 353 | if self.data_buffer_start >= self.data_buffer_size: |
| 354 | self.data_buffer_start -= self.data_buffer_size |
| 355 | self.data_buffer_end -= self.data_buffer_size |
| 356 | |
| 357 | monotonic_id_after = self.monotonic_id_start |
| 358 | # id wrap around |
| 359 | if monotonic_id_after >= monotonic_id_before: |
| 360 | return range(monotonic_id_before, monotonic_id_after) |
| 361 | else: |
| 362 | return chain( |
| 363 | range(monotonic_id_before, self.ID_MAX), range(0, monotonic_id_after) |
| 364 | ) |
| 365 | |
| 366 | |
| 367 | class ObjectSerde(ABC): |
| 368 | @abstractmethod |
| 369 | def serialize(self, value: Any) -> tuple[Any, int, bytes, int]: |
| 370 | """Serialize an object to bytes.""" |
| 371 | raise NotImplementedError |
| 372 | |
| 373 | @abstractmethod |
| 374 | def deserialize(self, data: memoryview) -> Any: |
| 375 | """Deserialize bytes back to an object.""" |
| 376 | raise NotImplementedError |
| 377 | |
| 378 | |
| 379 | class MsgpackSerde(ObjectSerde): |
| 380 | def __init__(self): |
| 381 | # Delayed import to avoid circular dependency |
| 382 | from vllm.multimodal.inputs import MultiModalKwargsItem |
| 383 | from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder |
| 384 | |
| 385 | self.encoder = MsgpackEncoder() |
| 386 | self.tensor_decoder = MsgpackDecoder(torch.Tensor, share_mem=False) |
| 387 | self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem, share_mem=False) |
| 388 | self._mm_kwargs_item_cls = MultiModalKwargsItem |
| 389 | |
| 390 | def serialize(self, value: Any) -> tuple[bytes | list[bytes], int, bytes, int]: |
| 391 | len_arr = None |
| 392 | if isinstance(value, (torch.Tensor, self._mm_kwargs_item_cls)): |
| 393 | type_name = type(value).__name__ |
| 394 | value = self.encoder.encode(value) |
| 395 | len_arr = [len(s) for s in value] |
| 396 | nbytes = sum(len_arr) |
| 397 | else: |
| 398 | value = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL) |
| 399 | type_name = type(value).__name__ |
| 400 | nbytes = len(value) |
| 401 | |
| 402 | object_metadata = (type_name, nbytes, len_arr) |
| 403 | serialized_metadata = pickle.dumps( |
| 404 | object_metadata, protocol=pickle.HIGHEST_PROTOCOL |
| 405 | ) |
| 406 | return value, nbytes, serialized_metadata, len(serialized_metadata) |
| 407 | |
| 408 | def deserialize(self, data_view: memoryview) -> Any: |
| 409 | # pickle.loads do not read past the end of a pickled object |
| 410 | # within a large buffer, so we can skip storing the metadata size |
| 411 | type_name, nbytes, len_arr = pickle.loads(data_view) |
| 412 | serialized_data = data_view[-nbytes:] |
| 413 | |
| 414 | if type_name == torch.Tensor.__name__: |
| 415 | obj = [] |
| 416 | start_idx = 0 |
| 417 | for length in len_arr: |
| 418 | item_bytes = serialized_data[start_idx : start_idx + length] |
| 419 | obj.append(item_bytes) |
| 420 | start_idx += length |
| 421 | obj = self.tensor_decoder.decode(obj) |
| 422 | elif type_name == self._mm_kwargs_item_cls.__name__: |
| 423 | obj = [] |
| 424 | start_idx = 0 |
| 425 | for length in len_arr: |
| 426 | item_bytes = serialized_data[start_idx : start_idx + length] |
| 427 | obj.append(item_bytes) |
| 428 | start_idx += length |
| 429 | obj = self.mm_decoder.decode(obj) |
| 430 | elif type_name == bytes.__name__: |
| 431 | obj = pickle.loads(serialized_data) |
| 432 | else: |
| 433 | raise ValueError(f"Unsupported object type '{type_name}' in metadata") |
| 434 | |
| 435 | return obj |
| 436 | |
| 437 | |
| 438 | @dataclass |
| 439 | class ShmObjectStorageHandle: |
| 440 | max_object_size: int |
| 441 | n_readers: int |
| 442 | ring_buffer_handle: tuple[int, str] |
| 443 | serde_class: type[ObjectSerde] |
| 444 | reader_lock: LockType | None |
| 445 | |
| 446 | |
| 447 | class SingleWriterShmObjectStorage: |
| 448 | """A single-writer, multiple-reader object storage system built on top of a |
| 449 | shared memory ring buffer. Provides key-value storage with automatic memory |
| 450 | management and cross-process serialization support. |
| 451 | |
| 452 | This storage system follows a FIFO (First-In-First-Out) eviction policy |
| 453 | where the oldest objects are automatically freed when memory runs low. |
| 454 | Memory is reclaimed based on reader reference counting - objects are only |
| 455 | freed when all readers have finished accessing them. |
| 456 | |
| 457 | Architecture: |
| 458 | - Single writer process can put(key, value) objects |
| 459 | - Multiple reader processes can get(address, monotonic_id, signature, key) |
| 460 | objects |
| 461 | - Built on SingleWriterShmRingBuffer for efficient shared memory management |
| 462 | - Thread-safe operations with reader synchronization via locks |
| 463 | |
| 464 | Key Features: |
| 465 | - FIFO Eviction: Oldest objects are evicted first when memory is full |
| 466 | - Reference Counting: Objects are only freed when no readers are |
| 467 | accessing them |
| 468 | - Duplicate Key Handling: Existing keys are not overwritten, just |
| 469 | re-referenced |
| 470 | - Customized Serialization: By default uses Msgpack for efficient |
| 471 | serialization of Python objects, but can be extended for custom types |
| 472 | - Cross-Process Safety: Uses shared memory with proper synchronization |
| 473 | - Automatic Cleanup: Garbage collection happens transparently during |
| 474 | allocation |
| 475 | |
| 476 | Handle Signatures: |
| 477 | - Handles are signed with an HMAC of (key, address, monotonic_id), keyed by |
| 478 | a random secret stored at the start of the shared memory |
| 479 | - The writer signs the handles it issues and rejects stale or forged |
| 480 | handles supplied by requests |
| 481 | - Readers verify the signature in get/touch before reading shared memory |
| 482 | |
| 483 | Memory Layout per Object: |
| 484 | `[4-byte reference_count][metadata_size][serialized_object_data]` |
| 485 | |
| 486 | Thread Safety: |
| 487 | - Writer operations (put, clear) are single-threaded by design |
| 488 | - Reader operations (get) are thread-safe with lock-based reference |
| 489 | counting |
| 490 | - Memory reclamation is handled exclusively by the writer process |
| 491 | """ |
| 492 | |
| 493 | def __init__( |
| 494 | self, |
| 495 | max_object_size: int, |
| 496 | n_readers: int, |
| 497 | ring_buffer: SingleWriterShmRingBuffer, |
| 498 | serde_class: type[ObjectSerde] = MsgpackSerde, |
| 499 | reader_lock: LockType | None = None, |
| 500 | ): |
| 501 | """Initialize the object storage. |
| 502 | |
| 503 | Args: |
| 504 | max_object_size: Maximum size for a single object in bytes. |
| 505 | n_readers: Number of reader processes that can access the storage. |
| 506 | ring_buffer: The shared memory ring buffer for storing objects. |
| 507 | serde_class: Serializer/deserializer for objects. |
| 508 | reader_lock: Optional lock for synchronizing reader access. |
| 509 | |
| 510 | Raises: |
| 511 | ValueError: If reader_lock is None for readers. |
| 512 | |
| 513 | """ |
| 514 | self.max_object_size = max_object_size |
| 515 | self.n_readers = n_readers |
| 516 | self.serde_class = serde_class |
| 517 | self.ser_de = serde_class() |
| 518 | self.ring_buffer = ring_buffer |
| 519 | self.is_writer = self.ring_buffer.is_writer |
| 520 | |
| 521 | self.flag_bytes = 4 # for in-use flag |
| 522 | |
| 523 | if self.is_writer: |
| 524 | # Key-value mapping: key -> (address, monotonic_id) |
| 525 | self.key_index: dict[str, tuple[int, int]] = {} |
| 526 | # Reverse mapping: monotonic_id -> key |
| 527 | self.id_index: dict[int, str] = {} |
| 528 | # Writer flag to track in-use status: monotonic_id -> count |
| 529 | self.writer_flag: dict[int, int] = {} |
| 530 | # Items touch() keeps from being freed until they are looked up |
| 531 | # or release_touches() is called |
| 532 | self._touched: set[int] = set() |
| 533 | else: |
| 534 | if reader_lock is None: |
| 535 | raise ValueError("Lock must be provided for readers.") |
| 536 | |
| 537 | self._reader_lock = reader_lock |
| 538 | |
| 539 | def clear(self) -> None: |
| 540 | """Clear the object storage.""" |
| 541 | if self.is_writer: |
| 542 | self.ring_buffer.clear() |
| 543 | self.key_index.clear() |
| 544 | self.id_index.clear() |
| 545 | self.writer_flag.clear() |
| 546 | self._touched.clear() |
| 547 | logger.debug("Object storage cleared and reinitialized.") |
| 548 | |
| 549 | def copy_to_buffer( |
| 550 | self, |
| 551 | data: bytes | list[bytes], |
| 552 | data_bytes: int, |
| 553 | metadata: bytes, |
| 554 | md_bytes: int, |
| 555 | data_view: memoryview, |
| 556 | ) -> None: |
| 557 | data_view[self.flag_bytes : self.flag_bytes + md_bytes] = metadata |
| 558 | if isinstance(data, bytes): |
| 559 | data_view[-data_bytes:] = data |
| 560 | elif isinstance(data, list): |
| 561 | start_idx = self.flag_bytes + md_bytes |
| 562 | for item_bytes in data: |
| 563 | item_size = len(item_bytes) |
| 564 | data_view[start_idx : start_idx + item_size] = item_bytes |
| 565 | start_idx += item_size |
| 566 | else: |
| 567 | raise ValueError(f"Unsupported data type for serialization: {type(data)}") |
| 568 | |
| 569 | def _validate_monotonic_id( |
| 570 | self, |
| 571 | address: int, |
| 572 | monotonic_id: int, |
| 573 | buf_metadata: tuple[int, int], |
| 574 | ) -> None: |
| 575 | if buf_metadata[0] != monotonic_id: |
| 576 | raise ValueError( |
| 577 | f"Data for address:id '{address}:{monotonic_id}'" |
| 578 | " has been modified or is invalid." |
| 579 | ) |
| 580 | |
| 581 | def _verify_reader_signature( |
| 582 | self, |
| 583 | address: int, |
| 584 | monotonic_id: int, |
| 585 | signature: list[int] | None, |
| 586 | key: str | None, |
| 587 | ) -> None: |
| 588 | if signature is None or key is None: |
| 589 | raise ValueError("Missing SHM handle signature for cache key.") |
| 590 | if ( |
| 591 | not isinstance(signature, list) |
| 592 | or len(signature) != hashlib.sha256().digest_size |
| 593 | or any(type(byte) is not int or not 0 <= byte <= 255 for byte in signature) |
| 594 | or not isinstance(key, str) |
| 595 | ): |
| 596 | raise ValueError("Invalid SHM handle signature for cache key.") |
| 597 | |
| 598 | expected_signature = self._make_signature(key, address, monotonic_id) |
| 599 | if not hmac.compare_digest(bytes(signature), expected_signature): |
| 600 | raise ValueError("Invalid SHM handle signature for cache key.") |
| 601 | |
| 602 | def verify_signature( |
| 603 | self, |
| 604 | key: str | None, |
| 605 | address: int, |
| 606 | monotonic_id: int, |
| 607 | signature: list[int] | None, |
| 608 | ) -> None: |
| 609 | """Verify that a handle was issued for the given key. |
| 610 | |
| 611 | For writers: the handle must also be the key's current entry |
| 612 | For readers: only the signature is checked |
| 613 | |
| 614 | Args: |
| 615 | key: String key the handle belongs to |
| 616 | address: Address of the object |
| 617 | monotonic_id: Monotonic ID of the object |
| 618 | signature: Signature issued with the handle |
| 619 | |
| 620 | """ |
| 621 | if self.is_writer and ( |
| 622 | key is None or self.key_index.get(key) != (address, monotonic_id) |
| 623 | ): |
| 624 | raise ValueError("Invalid SHM handle signature for cache key.") |
| 625 | self._verify_reader_signature(address, monotonic_id, signature, key) |
| 626 | |
| 627 | def _make_signature(self, key: str, address: int, monotonic_id: int) -> bytes: |
| 628 | payload = ( |
| 629 | hashlib.sha256(key.encode("utf-8")).digest() |
| 630 | + address.to_bytes(8, "little", signed=False) |
| 631 | + monotonic_id.to_bytes(4, "little", signed=False) |
| 632 | ) |
| 633 | return hmac.new( |
| 634 | self.ring_buffer.get_auth_secret(), |
| 635 | payload, |
| 636 | hashlib.sha256, |
| 637 | ).digest() |
| 638 | |
| 639 | def increment_writer_flag(self, id: int) -> None: |
| 640 | """Set the in-use flag for the writer.""" |
| 641 | self.writer_flag[id] = self.writer_flag.get(id, 0) + 1 |
| 642 | |
| 643 | def increment_reader_flag(self, data_view: memoryview) -> None: |
| 644 | """Set the in-use flag for the reader.""" |
| 645 | # >0 for in-use flag |
| 646 | reader_count = self.ring_buffer.byte2int(data_view) |
| 647 | data_view[:] = self.ring_buffer.int2byte(reader_count + 1) |
| 648 | |
| 649 | def free_unused(self) -> None: |
| 650 | """Free unused buffers in the ring buffer.""" |
| 651 | # try to free up 2*max_object_size bytes of space in the ring buffer, |
| 652 | # since the buffer might be fragmented |
| 653 | freed_ids = self.ring_buffer.free_buf( |
| 654 | self.default_is_free_check, 2 * self.max_object_size |
| 655 | ) |
| 656 | # update the metadata after freeing up space |
| 657 | for freed_id in freed_ids: |
| 658 | key_to_free = self.id_index[freed_id] |
| 659 | del self.key_index[key_to_free] |
| 660 | del self.id_index[freed_id] |
| 661 | del self.writer_flag[freed_id] |
| 662 | |
| 663 | def is_cached(self, key: str) -> bool: |
| 664 | """Check if the object with the given key is cached.""" |
| 665 | return key in self.key_index |
| 666 | |
| 667 | def get_cached(self, key: str) -> tuple[int, int]: |
| 668 | """Get the cached object by key if it exists.""" |
| 669 | address, monotonic_id = self.key_index[key] |
| 670 | self.increment_writer_flag(monotonic_id) |
| 671 | self._touched.discard(monotonic_id) |
| 672 | return address, monotonic_id |
| 673 | |
| 674 | def release_touches(self) -> None: |
| 675 | """Stop protecting touched items that were not looked up.""" |
| 676 | self._touched.clear() |
| 677 | |
| 678 | def get_signature(self, key: str) -> list[int]: |
| 679 | """Sign the handle of a cached object so readers can verify it.""" |
| 680 | address, monotonic_id = self.key_index[key] |
| 681 | return list(self._make_signature(key, address, monotonic_id)) |
| 682 | |
| 683 | def put(self, key: str, value: Any) -> tuple[int, int]: |
| 684 | """Store a key-value pair in the object storage. |
| 685 | Attempts to free max_object_size bytes using FIFO order |
| 686 | when the ring buffer runs out of space during a put() operation. |
| 687 | |
| 688 | Args: |
| 689 | key: String key to identify the object |
| 690 | value: Any serializable Python object |
| 691 | |
| 692 | Raises: |
| 693 | MemoryError: If there's not enough space in the buffer |
| 694 | ValueError: If the serialized object is too large |
| 695 | ValueError: If the key already exists in the storage |
| 696 | |
| 697 | """ |
| 698 | if key in self.key_index: |
| 699 | raise ValueError(f"Key '{key}' already exists in the storage.") |
| 700 | |
| 701 | object_data, data_bytes, object_metadata, md_bytes = self.ser_de.serialize( |
| 702 | value |
| 703 | ) |
| 704 | buffer_size = self.flag_bytes + data_bytes + md_bytes |
| 705 | # Sanity checks |
| 706 | if buffer_size > self.max_object_size: |
| 707 | raise ValueError( |
| 708 | f"Serialized object size ({buffer_size} bytes) exceeds " |
| 709 | f"max object size ({self.max_object_size} bytes)" |
| 710 | ) |
| 711 | |
| 712 | # Allocate new buffer |
| 713 | try: |
| 714 | address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size) |
| 715 | except MemoryError: |
| 716 | self.free_unused() |
| 717 | # try again after freeing up space |
| 718 | address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size) |
| 719 | |
| 720 | # Write data to buffer |
| 721 | with self.ring_buffer.access_buf(address) as (data_view, metadata): |
| 722 | data_view[: self.flag_bytes] = self.ring_buffer.int2byte(0) |
| 723 | self.copy_to_buffer( |
| 724 | object_data, data_bytes, object_metadata, md_bytes, data_view |
| 725 | ) |
| 726 | self.increment_writer_flag(monotonic_id) |
| 727 | |
| 728 | # Update key index |
| 729 | self.key_index[key] = (address, monotonic_id) |
| 730 | self.id_index[monotonic_id] = key |
| 731 | return address, monotonic_id |
| 732 | |
| 733 | def get( |
| 734 | self, |
| 735 | address: int, |
| 736 | monotonic_id: int, |
| 737 | signature: list[int] | None = None, |
| 738 | key: str | None = None, |
| 739 | ) -> Any: |
| 740 | if not self.is_writer: |
| 741 | # Reject forged handles before dereferencing the supplied address. |
| 742 | self.verify_signature(key, address, monotonic_id, signature) |
| 743 | |
| 744 | # Read data from buffer |
| 745 | with self.ring_buffer.access_buf(address) as (data_view, buf_metadata): |
| 746 | self._validate_monotonic_id(address, monotonic_id, buf_metadata) |
| 747 | |
| 748 | obj = self.ser_de.deserialize(data_view[self.flag_bytes :]) |
| 749 | |
| 750 | # decrease the in-use flag for reader reads |
| 751 | if self._reader_lock is not None: |
| 752 | with self._reader_lock: |
| 753 | self.increment_reader_flag(data_view[: self.flag_bytes]) |
| 754 | else: |
| 755 | # if self._reader_lock is None, it means we are the writer |
| 756 | # in this case, we do not need to decrease the reader count |
| 757 | assert self.is_writer |
| 758 | |
| 759 | return obj |
| 760 | |
| 761 | def touch( |
| 762 | self, |
| 763 | key: str, |
| 764 | address: int = 0, |
| 765 | monotonic_id: int = 0, |
| 766 | signature: list[int] | None = None, |
| 767 | ) -> None: |
| 768 | """Touch an existing cached item. |
| 769 | |
| 770 | For writers (ShmObjectStoreSenderCache): Protect the item from |
| 771 | eviction until get_cached or release_touches |
| 772 | For readers (ShmObjectStoreReceiverCache): Validate the handle |
| 773 | |
| 774 | Args: |
| 775 | key: String key of the object to touch |
| 776 | address: Address of the object (only for readers) |
| 777 | monotonic_id: Monotonic ID of the object (only for readers) |
| 778 | signature: Server-issued handle signature (only for readers) |
| 779 | |
| 780 | """ |
| 781 | if self._reader_lock is None: |
| 782 | if key not in self.key_index: |
| 783 | return None |
| 784 | self._touched.add(self.key_index[key][1]) |
| 785 | else: |
| 786 | # Reject forged handles before dereferencing the supplied address. |
| 787 | self.verify_signature(key, address, monotonic_id, signature) |
| 788 | with self.ring_buffer.access_buf(address) as (_, buf_metadata): |
| 789 | self._validate_monotonic_id(address, monotonic_id, buf_metadata) |
| 790 | |
| 791 | def close(self) -> None: |
| 792 | """Close the shared memory.""" |
| 793 | self.ring_buffer.close() |
| 794 | |
| 795 | def handle(self): |
| 796 | """Get handle for sharing across processes.""" |
| 797 | return ShmObjectStorageHandle( |
| 798 | max_object_size=self.max_object_size, |
| 799 | n_readers=self.n_readers, |
| 800 | ring_buffer_handle=self.ring_buffer.handle(), |
| 801 | serde_class=self.serde_class, |
| 802 | reader_lock=self._reader_lock, |
| 803 | ) |
| 804 | |
| 805 | @staticmethod |
| 806 | def create_from_handle( |
| 807 | handle: ShmObjectStorageHandle, |
| 808 | ) -> "SingleWriterShmObjectStorage": |
| 809 | logger.debug("Creating storage from handle: %s", handle) |
| 810 | ring_buffer = SingleWriterShmRingBuffer(*handle.ring_buffer_handle) |
| 811 | return SingleWriterShmObjectStorage( |
| 812 | max_object_size=handle.max_object_size, |
| 813 | n_readers=handle.n_readers, |
| 814 | ring_buffer=ring_buffer, |
| 815 | serde_class=handle.serde_class, |
| 816 | reader_lock=handle.reader_lock, |
| 817 | ) |
| 818 | |
| 819 | def default_is_free_check(self, id: int, buf: memoryview) -> bool: |
| 820 | """Default is_free function that checks if the first 4 bytes are zero. |
| 821 | This indicates that the buffer is free. |
| 822 | """ |
| 823 | if id in self._touched: |
| 824 | return False |
| 825 | reader_count = int.from_bytes(buf[0:4], "little", signed=True) |
| 826 | writer_count = self.writer_flag[id] |
| 827 | return reader_count >= writer_count * self.n_readers |