| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import contextlib |
| 5 | import logging |
| 6 | import mmap |
| 7 | import os |
| 8 | import random |
| 9 | import threading |
| 10 | |
| 11 | try: |
| 12 | from vllm.fs_io_C import ( # pyright: ignore[reportMissingImports] |
| 13 | batch_load_block as batch_load_block_C, |
| 14 | ) |
| 15 | from vllm.fs_io_C import ( |
| 16 | batch_store_block as batch_store_block_C, |
| 17 | ) |
| 18 | |
| 19 | _HAS_FSIO_C = True |
| 20 | except ImportError: |
| 21 | _HAS_FSIO_C = False |
| 22 | |
| 23 | logger = logging.getLogger(__name__) |
| 24 | |
| 25 | # O_DIRECT is Linux-specific and not available on macOS |
| 26 | O_DIRECT = getattr(os, "O_DIRECT", 0) |
| 27 | |
| 28 | # Thread-local storage for unique temporary file suffixes |
| 29 | _thread_local = threading.local() |
| 30 | |
| 31 | |
| 32 | def _get_tmp_suffix() -> str: |
| 33 | """Generate a thread-local unique suffix for temporary files.""" |
| 34 | try: |
| 35 | return _thread_local.tmp_suffix |
| 36 | except AttributeError: |
| 37 | _thread_local.tmp_suffix = f"_{random.randint(0, 2**63 - 1)}.tmp" |
| 38 | return _thread_local.tmp_suffix |
| 39 | |
| 40 | |
| 41 | def probe_o_direct(directory: str) -> bool: |
| 42 | """Return whether ``O_DIRECT`` I/O works in *directory*. |
| 43 | |
| 44 | ``O_DIRECT`` is unsupported on some filesystems (e.g. the overlayfs backing |
| 45 | a container ``/tmp``, older tmpfs, or some NFS mounts), where opening or |
| 46 | writing a file with it fails with ``EINVAL``. Probe once with an aligned |
| 47 | single-page write so callers can fall back to buffered I/O instead of |
| 48 | failing on every block. |
| 49 | """ |
| 50 | if not O_DIRECT: |
| 51 | return False |
| 52 | path = os.path.join(directory, f".o_direct_probe{_get_tmp_suffix()}") |
| 53 | page = mmap.mmap(-1, mmap.PAGESIZE) |
| 54 | try: |
| 55 | fd = os.open(path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC | O_DIRECT, 0o644) |
| 56 | try: |
| 57 | os.write(fd, page) |
| 58 | finally: |
| 59 | os.close(fd) |
| 60 | return True |
| 61 | except OSError: |
| 62 | return False |
| 63 | finally: |
| 64 | page.close() |
| 65 | with contextlib.suppress(OSError): |
| 66 | os.remove(path) |
| 67 | |
| 68 | |
| 69 | def _ensure_dirs(path: str) -> None: |
| 70 | """Create parent directories of *path* if they don't exist.""" |
| 71 | os.makedirs(os.path.dirname(path), exist_ok=True) |
| 72 | |
| 73 | |
| 74 | def _validate_offsets(view: memoryview, offsets: list[int], block_size: int) -> None: |
| 75 | """Raise if any block would read/write past the bounds of `view`. |
| 76 | |
| 77 | Without this, an out-of-range offset silently clips to a shorter (or |
| 78 | empty) slice instead of failing, since memoryview slicing follows |
| 79 | Python's slice-clamping semantics rather than raising. |
| 80 | """ |
| 81 | total_len = len(view.cast("B")) |
| 82 | for offset in offsets: |
| 83 | if offset < 0 or offset + block_size > total_len: |
| 84 | raise ValueError( |
| 85 | f"block offset {offset} (block_size {block_size}) is out of " |
| 86 | f"bounds for a buffer of size {total_len}" |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | def _store_block( |
| 91 | dest_path: str, |
| 92 | buffer: memoryview, |
| 93 | offset: int, |
| 94 | block_size: int, |
| 95 | use_o_direct: bool = True, |
| 96 | ) -> None: |
| 97 | """Store callback: write to a temp file then atomically replace the target.""" |
| 98 | # Check if block already exists to avoid redundant writes |
| 99 | if os.path.exists(dest_path): |
| 100 | return |
| 101 | |
| 102 | tmp_path = dest_path + _get_tmp_suffix() |
| 103 | # Ensure parent directories exist |
| 104 | _ensure_dirs(dest_path) |
| 105 | |
| 106 | # Write block atomically. Cast to a flat byte view so the slice uses byte |
| 107 | # indices; the raw memoryview may be multi-dimensional with itemsize > 1. |
| 108 | view_slice = buffer.cast("B")[offset : offset + block_size] |
| 109 | o_direct = O_DIRECT if use_o_direct else 0 |
| 110 | try: |
| 111 | fd = os.open( |
| 112 | tmp_path, |
| 113 | os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | o_direct, |
| 114 | 0o644, |
| 115 | ) |
| 116 | try: |
| 117 | written = os.write(fd, view_slice) |
| 118 | if written < len(view_slice): |
| 119 | raise OSError( |
| 120 | f"Short write: expected {len(view_slice)} bytes, wrote {written}" |
| 121 | ) |
| 122 | finally: |
| 123 | os.close(fd) |
| 124 | os.replace(tmp_path, dest_path) |
| 125 | except Exception: |
| 126 | try: |
| 127 | os.remove(tmp_path) |
| 128 | except OSError as cleanup_exc: |
| 129 | logger.warning("Failed to remove temp file %s: %s", tmp_path, cleanup_exc) |
| 130 | raise |
| 131 | |
| 132 | |
| 133 | def _load_block( |
| 134 | source_path: str, |
| 135 | view: memoryview, |
| 136 | offset: int, |
| 137 | block_size: int, |
| 138 | use_o_direct: bool = True, |
| 139 | ) -> None: |
| 140 | """Read one KV block from disk; remove the file only on a provable short |
| 141 | read (a too-short file is genuine corruption) and leave it untouched on any |
| 142 | other error.""" |
| 143 | fd: int | None = None |
| 144 | view_slice = view.cast("B")[offset : offset + block_size] |
| 145 | o_direct = O_DIRECT if use_o_direct else 0 |
| 146 | |
| 147 | try: |
| 148 | fd = os.open(source_path, os.O_RDONLY | o_direct) |
| 149 | bytes_read = os.readv(fd, [view_slice]) |
| 150 | if bytes_read < block_size: |
| 151 | # A failure to remove must not mask the short-read error below. |
| 152 | try: |
| 153 | os.remove(source_path) |
| 154 | except OSError as cleanup_exc: |
| 155 | logger.warning( |
| 156 | "Failed to remove short-read file %s: %s", |
| 157 | source_path, |
| 158 | cleanup_exc, |
| 159 | ) |
| 160 | raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}") |
| 161 | finally: |
| 162 | if fd is not None: |
| 163 | os.close(fd) |
| 164 | |
| 165 | |
| 166 | def batch_store_block( |
| 167 | paths: list[str], |
| 168 | view: memoryview, |
| 169 | offsets: list[int], |
| 170 | block_size: int, |
| 171 | use_o_direct: bool = True, |
| 172 | ) -> None: |
| 173 | """Store a batch of KV blocks from a shared buffer to disk in one call. |
| 174 | |
| 175 | Each block buffer[offsets[i] : offsets[i]+block_size] is written atomically |
| 176 | to dest_paths[i] via a temp-file rename. Raises on first error. |
| 177 | """ |
| 178 | _validate_offsets(view, offsets, block_size) |
| 179 | |
| 180 | if _HAS_FSIO_C: |
| 181 | view_B = view.cast("B") |
| 182 | view_slices = [view_B[x : x + block_size] for x in offsets] |
| 183 | tmp_paths = [p + _get_tmp_suffix() for p in paths] |
| 184 | return batch_store_block_C(tmp_paths, paths, view_slices, use_o_direct) |
| 185 | else: |
| 186 | for path, offset in zip(paths, offsets): |
| 187 | _store_block(path, view, offset, block_size, use_o_direct) |
| 188 | |
| 189 | |
| 190 | def batch_load_block( |
| 191 | paths: list[str], |
| 192 | view: memoryview, |
| 193 | offsets: list[int], |
| 194 | block_size: int, |
| 195 | use_o_direct: bool = True, |
| 196 | ) -> None: |
| 197 | """Load a batch of KV blocks from disk into a shared buffer in one call. |
| 198 | |
| 199 | Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size]. |
| 200 | Raises on first error (see _load_block for the delete-on-short-read policy). |
| 201 | On failure the raised OSError carries ``num_succeeded`` = the number of |
| 202 | blocks loaded before the failing one, so the tier can keep them. |
| 203 | """ |
| 204 | _validate_offsets(view, offsets, block_size) |
| 205 | |
| 206 | if _HAS_FSIO_C: |
| 207 | view_B = view.cast("B") |
| 208 | view_slices = [view_B[x : x + block_size] for x in offsets] |
| 209 | return batch_load_block_C(paths, view_slices, use_o_direct) |
| 210 | else: |
| 211 | for i, (path, offset) in enumerate(zip(paths, offsets)): |
| 212 | try: |
| 213 | _load_block(path, view, offset, block_size, use_o_direct) |
| 214 | except OSError as exc: |
| 215 | # Blocks 0..i-1 loaded fine; record the count for partial keep. |
| 216 | # The C path sets the same attribute via PyObject_SetAttrString. |
| 217 | exc.num_succeeded = i # type: ignore[attr-defined] |
| 218 | raise |