| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """WeightCacheKey fingerprinting and socket protocol for the weight cache daemon. |
| 4 | |
| 5 | The protocol uses pickle over a Unix domain socket and is only intended for |
| 6 | communication between trusted local processes owned by the same user. The |
| 7 | sockets live in a per-user private directory (mode 0700) and the daemon |
| 8 | restricts the socket file permissions to the owner (0600). Both the daemon and |
| 9 | the engine verify that the directory and socket are owned by the current user |
| 10 | and are not group/world accessible before trusting them, so a different local |
| 11 | user cannot pre-plant a malicious socket at a predictable path. |
| 12 | """ |
| 13 | |
| 14 | import glob |
| 15 | import json |
| 16 | import os |
| 17 | import pickle |
| 18 | import socket |
| 19 | import stat |
| 20 | import struct |
| 21 | import tempfile |
| 22 | from dataclasses import dataclass, fields |
| 23 | from typing import Any, NamedTuple |
| 24 | |
| 25 | import torch |
| 26 | from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor |
| 27 | from transformers.utils import SAFE_WEIGHTS_INDEX_NAME |
| 28 | |
| 29 | import vllm.version |
| 30 | from vllm.config import ModelConfig |
| 31 | from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase |
| 32 | from vllm.model_executor.model_loader.weight_cache.utils import ( |
| 33 | format_socket_role_suffix, |
| 34 | ) |
| 35 | from vllm.model_executor.model_loader.weight_utils import ( |
| 36 | filter_duplicate_safetensors_files, |
| 37 | ) |
| 38 | from vllm.platforms import current_platform |
| 39 | from vllm.utils.hashing import safe_hash |
| 40 | |
| 41 | SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}{role}.sock" |
| 42 | SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}" |
| 43 | |
| 44 | _LEN_STRUCT = struct.Struct("!Q") |
| 45 | # Sanity bound for a single message. IPC handles are tiny; only small |
| 46 | # non-CUDA tensors are ever shipped by value. |
| 47 | MAX_MSG_SIZE = 1 << 34 |
| 48 | |
| 49 | |
| 50 | def _current_uid() -> int: |
| 51 | getuid = getattr(os, "getuid", None) |
| 52 | return getuid() if getuid is not None else -1 |
| 53 | |
| 54 | |
| 55 | class WeightCacheUnavailableError(Exception): |
| 56 | """Raised when no weight cache daemon is reachable or usable.""" |
| 57 | |
| 58 | |
| 59 | class CacheConfigMismatchError(Exception): |
| 60 | """Raised when the daemon's cached weights don't match the engine.""" |
| 61 | |
| 62 | |
| 63 | class UnsupportedQuantForIPCError(Exception): |
| 64 | """Raised when a quantization method is not verified for IPC weight sharing.""" |
| 65 | |
| 66 | |
| 67 | class UnsupportedPlatformForIPCError(Exception): |
| 68 | """Raised when the current platform cannot share CUDA IPC handles.""" |
| 69 | |
| 70 | |
| 71 | def check_ipc_platform_support() -> None: |
| 72 | """Hard-error unless the current platform can share CUDA IPC handles. |
| 73 | |
| 74 | Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other |
| 75 | platforms (e.g. XPU) would silently ship every tensor by value instead. |
| 76 | |
| 77 | Raises: |
| 78 | UnsupportedPlatformForIPCError: If the current platform is not |
| 79 | CUDA/ROCm. |
| 80 | |
| 81 | """ |
| 82 | if current_platform.is_cuda_alike(): |
| 83 | return |
| 84 | raise UnsupportedPlatformForIPCError( |
| 85 | f"platform {current_platform.device_name!r} does not support CUDA IPC " |
| 86 | "weight sharing; only CUDA and ROCm are supported. Use the default " |
| 87 | "--load-format for this platform." |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | # The daemon transfers post processed weights directly |
| 92 | def check_ipc_quant_support(model: torch.nn.Module) -> None: |
| 93 | """Hard-error unless every quant method supports pre-processed weights. |
| 94 | |
| 95 | Args: |
| 96 | model: The model to inspect (weights need not be loaded). |
| 97 | |
| 98 | Raises: |
| 99 | UnsupportedQuantForIPCError: If any quant method does not declare |
| 100 | ``supports_pre_processed_weights``. |
| 101 | |
| 102 | """ |
| 103 | for name, module in model.named_modules(): |
| 104 | quant_method = getattr(module, "quant_method", None) |
| 105 | if ( |
| 106 | isinstance(quant_method, QuantizeMethodBase) |
| 107 | and not quant_method.supports_pre_processed_weights |
| 108 | ): |
| 109 | raise UnsupportedQuantForIPCError( |
| 110 | f"layer {name or '<root>'}: {type(quant_method).__name__} " |
| 111 | "does not support loading from pre-processed weights." |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | def get_current_device_uuid() -> str: |
| 116 | """UUID of the physical GPU backing the current accelerator device.""" |
| 117 | return current_platform.get_device_uuid(torch.accelerator.current_device_index()) |
| 118 | |
| 119 | |
| 120 | def get_socket_dir(socket_dir: str | None = None) -> str: |
| 121 | """Return the directory that holds the daemon sockets. |
| 122 | |
| 123 | When no explicit directory is given, use a per-user private directory |
| 124 | under the system temp dir so its path is unpredictable to other users and |
| 125 | can be locked down to mode 0700. |
| 126 | """ |
| 127 | if socket_dir is not None: |
| 128 | return socket_dir |
| 129 | return os.path.join( |
| 130 | tempfile.gettempdir(), SOCKET_DIR_TEMPLATE.format(uid=_current_uid()) |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | def get_socket_path( |
| 135 | gpu_uuid: str, |
| 136 | socket_dir: str | None = None, |
| 137 | *, |
| 138 | is_draft: bool = False, |
| 139 | ) -> str: |
| 140 | """Socket path of a daemon group; ``is_draft=False`` is the target. |
| 141 | |
| 142 | The GPU uuid is hashed to keep the name well under the AF_UNIX path |
| 143 | limit (~108 bytes) even with the draft role suffix. |
| 144 | """ |
| 145 | gpu_id = safe_hash(gpu_uuid.encode()).hexdigest() |
| 146 | name = SOCKET_NAME_TEMPLATE.format( |
| 147 | gpu_uuid=gpu_id, |
| 148 | role=format_socket_role_suffix(is_draft), |
| 149 | ) |
| 150 | return os.path.join(get_socket_dir(socket_dir), name) |
| 151 | |
| 152 | |
| 153 | def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None: |
| 154 | """Create the socket directory (if needed) locked down to the owner. |
| 155 | |
| 156 | Called by the daemon before binding. Existing directories are re-checked |
| 157 | and, for the auto-derived path, tightened so a pre-existing world-writable |
| 158 | directory is rejected. |
| 159 | """ |
| 160 | os.makedirs(directory, mode=0o700, exist_ok=True) |
| 161 | if strict_perms: |
| 162 | os.chmod(directory, 0o700) |
| 163 | verify_private_dir(directory, strict_perms=strict_perms) |
| 164 | |
| 165 | |
| 166 | def verify_private_dir(directory: str, strict_perms: bool = True) -> None: |
| 167 | """Verify a directory is a real dir owned by us and not world/group readable. |
| 168 | |
| 169 | When ``strict_perms`` is False the group/world permission bits are not |
| 170 | checked; this is used for directories the operator explicitly configured |
| 171 | (they own the trust decision), while the auto-derived per-user directory is |
| 172 | always checked strictly. |
| 173 | """ |
| 174 | info = os.lstat(directory) |
| 175 | if stat.S_ISLNK(info.st_mode): |
| 176 | raise WeightCacheUnavailableError( |
| 177 | f"Refusing to use symlinked socket directory {directory}" |
| 178 | ) |
| 179 | if not stat.S_ISDIR(info.st_mode): |
| 180 | raise WeightCacheUnavailableError(f"{directory} is not a directory") |
| 181 | uid = _current_uid() |
| 182 | if uid != -1 and info.st_uid != uid: |
| 183 | raise WeightCacheUnavailableError( |
| 184 | f"Socket directory {directory} is not owned by the current user" |
| 185 | ) |
| 186 | if strict_perms and info.st_mode & 0o077: |
| 187 | raise WeightCacheUnavailableError( |
| 188 | f"Socket directory {directory} is group/world accessible" |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | def verify_socket_owner(socket_path: str, strict_perms: bool = True) -> None: |
| 193 | """Verify the socket lives in a private dir and is owned by the current user. |
| 194 | |
| 195 | Called by the engine before connecting so it never talks to a socket a |
| 196 | different user could have planted. |
| 197 | """ |
| 198 | verify_private_dir(os.path.dirname(socket_path), strict_perms=strict_perms) |
| 199 | info = os.lstat(socket_path) |
| 200 | if stat.S_ISLNK(info.st_mode): |
| 201 | raise WeightCacheUnavailableError( |
| 202 | f"Refusing to connect to symlinked socket {socket_path}" |
| 203 | ) |
| 204 | uid = _current_uid() |
| 205 | if uid != -1 and info.st_uid != uid: |
| 206 | raise WeightCacheUnavailableError( |
| 207 | f"Socket {socket_path} is not owned by the current user" |
| 208 | ) |
| 209 | |
| 210 | |
| 211 | def verify_peer_is_owner(conn: socket.socket) -> None: |
| 212 | """Best-effort check that the connecting peer runs as the current user. |
| 213 | |
| 214 | Uses SO_PEERCRED where available (Linux). Silently returns on platforms |
| 215 | that do not expose peer credentials. |
| 216 | """ |
| 217 | so_peercred = getattr(socket, "SO_PEERCRED", None) |
| 218 | if so_peercred is None: |
| 219 | return |
| 220 | try: |
| 221 | creds = conn.getsockopt(socket.SOL_SOCKET, so_peercred, struct.calcsize("3i")) |
| 222 | _, peer_uid, _ = struct.unpack("3i", creds) |
| 223 | except OSError: |
| 224 | return |
| 225 | uid = _current_uid() |
| 226 | if uid != -1 and peer_uid != uid: |
| 227 | raise PermissionError(f"Rejecting weight cache connection from uid {peer_uid}") |
| 228 | |
| 229 | |
| 230 | def _hash_quant_config(quant_config: Any) -> str: |
| 231 | if quant_config is None: |
| 232 | return "" |
| 233 | if hasattr(quant_config, "to_dict"): |
| 234 | quant_config = quant_config.to_dict() |
| 235 | payload = json.dumps(quant_config, sort_keys=True, default=str) |
| 236 | return safe_hash(payload.encode(), usedforsecurity=False).hexdigest() |
| 237 | |
| 238 | |
| 239 | def _safetensors_header(path: str) -> bytes: |
| 240 | """Return the raw safetensors header (length prefix + JSON) of a file. |
| 241 | |
| 242 | The header carries tensor names, dtypes, shapes and byte offsets, so it is |
| 243 | a content fingerprint of the shard without reading any weight bytes. |
| 244 | """ |
| 245 | with open(path, "rb") as f: |
| 246 | size_bytes = f.read(8) |
| 247 | (header_len,) = struct.unpack("<Q", size_bytes) |
| 248 | return size_bytes + f.read(header_len) |
| 249 | |
| 250 | |
| 251 | def hash_checkpoint(model: str) -> str | None: |
| 252 | """Fingerprint checkpoint content from local safetensors metadata. |
| 253 | |
| 254 | Hashes each shard's safetensors header so a daemon and an engine pointing |
| 255 | at identical weights in different directories produce the same key. Returns |
| 256 | None when local safetensors files can't be located (e.g. an undownloaded |
| 257 | Hugging Face repo id), leaving the caller to fall back to the model path. |
| 258 | """ |
| 259 | if not os.path.isdir(model): |
| 260 | return None |
| 261 | files = glob.glob(os.path.join(model, "*.safetensors")) |
| 262 | if os.path.isfile(os.path.join(model, SAFE_WEIGHTS_INDEX_NAME)): |
| 263 | files = filter_duplicate_safetensors_files( |
| 264 | files, model, SAFE_WEIGHTS_INDEX_NAME |
| 265 | ) |
| 266 | if not files: |
| 267 | return None |
| 268 | hasher = safe_hash(b"", usedforsecurity=False) |
| 269 | for path in sorted(files, key=os.path.basename): |
| 270 | hasher.update(os.path.basename(path).encode()) |
| 271 | hasher.update(_safetensors_header(path)) |
| 272 | return hasher.hexdigest() |
| 273 | |
| 274 | |
| 275 | @dataclass(frozen=True) |
| 276 | class WeightCacheKey: |
| 277 | """Fingerprint of the cached weights. |
| 278 | |
| 279 | Any mismatch between the daemon's and the engine's fingerprint means the |
| 280 | cached weights cannot be reused and the engine must load from disk. |
| 281 | """ |
| 282 | |
| 283 | checkpoint: str |
| 284 | model_arch: str |
| 285 | tp_size: int |
| 286 | tp_rank: int |
| 287 | dtype: str |
| 288 | quantization: str | None |
| 289 | quant_config_hash: str |
| 290 | revision: str | None |
| 291 | vllm_version: str |
| 292 | pp_size: int = 1 |
| 293 | pp_rank: int = 0 |
| 294 | is_draft: bool = False |
| 295 | """Daemon group the weights come from; False is the target model.""" |
| 296 | dp_size: int = 1 |
| 297 | dp_rank: int = 0 |
| 298 | |
| 299 | @classmethod |
| 300 | def from_model_config( |
| 301 | cls, |
| 302 | model_config: ModelConfig, |
| 303 | tp_size: int, |
| 304 | tp_rank: int, |
| 305 | *, |
| 306 | pp_size: int = 1, |
| 307 | pp_rank: int = 0, |
| 308 | is_draft: bool = False, |
| 309 | dp_size: int = 1, |
| 310 | dp_rank: int = 0, |
| 311 | ) -> "WeightCacheKey": |
| 312 | """Build the fingerprint for a model configuration. |
| 313 | |
| 314 | Must be called before weight loading: process_weights_after_loading |
| 315 | may mutate hf_config.quantization_config, which would change the hash |
| 316 | between the daemon and the engine. |
| 317 | |
| 318 | The checkpoint is identified by a hash of its safetensors metadata when |
| 319 | the weights are available locally, so a daemon and engine referencing |
| 320 | identical weights in different directories still match; otherwise it |
| 321 | falls back to the model path. |
| 322 | """ |
| 323 | hf_config = model_config.hf_config |
| 324 | arch = ",".join(getattr(hf_config, "architectures", None) or []) |
| 325 | quant_config = getattr(hf_config, "quantization_config", None) |
| 326 | checkpoint = hash_checkpoint(model_config.model) or model_config.model |
| 327 | return cls( |
| 328 | checkpoint=checkpoint, |
| 329 | model_arch=arch, |
| 330 | tp_size=tp_size, |
| 331 | tp_rank=tp_rank, |
| 332 | pp_size=pp_size, |
| 333 | pp_rank=pp_rank, |
| 334 | dtype=str(model_config.dtype), |
| 335 | quantization=model_config.quantization, |
| 336 | quant_config_hash=_hash_quant_config(quant_config), |
| 337 | revision=model_config.revision, |
| 338 | vllm_version=vllm.version.__version__, |
| 339 | is_draft=is_draft, |
| 340 | dp_size=dp_size, |
| 341 | dp_rank=dp_rank, |
| 342 | ) |
| 343 | |
| 344 | def mismatched_fields(self, other: "WeightCacheKey") -> list[str]: |
| 345 | return [ |
| 346 | f.name |
| 347 | for f in fields(self) |
| 348 | if getattr(self, f.name) != getattr(other, f.name) |
| 349 | ] |
| 350 | |
| 351 | |
| 352 | @dataclass |
| 353 | class TensorEntry: |
| 354 | """A single cached tensor. |
| 355 | |
| 356 | CUDA tensors are exported as `torch.multiprocessing` reduction args |
| 357 | (CUDA IPC handles); non-CUDA tensors are shipped by value. |
| 358 | """ |
| 359 | |
| 360 | kind: str |
| 361 | """Either "param" or "buffer".""" |
| 362 | ipc_args: tuple | None = None |
| 363 | cpu_tensor: torch.Tensor | None = None |
| 364 | |
| 365 | @classmethod |
| 366 | def from_tensor(cls, tensor: torch.Tensor, kind: str) -> "TensorEntry": |
| 367 | tensor = tensor.detach() |
| 368 | if tensor.is_cuda: |
| 369 | _, ipc_args = reduce_tensor(tensor) |
| 370 | return cls(kind=kind, ipc_args=ipc_args) |
| 371 | return cls(kind=kind, cpu_tensor=tensor.cpu()) |
| 372 | |
| 373 | def rebuild(self, device_index: int) -> torch.Tensor: |
| 374 | if self.ipc_args is None: |
| 375 | assert self.cpu_tensor is not None |
| 376 | return self.cpu_tensor |
| 377 | args = list(self.ipc_args) |
| 378 | # Index 6 of the args from reduce_tensor is the device index. It must |
| 379 | # be retargeted to the local index since the daemon and the engine may |
| 380 | # have different CUDA_VISIBLE_DEVICES mappings. |
| 381 | args[6] = device_index |
| 382 | return rebuild_cuda_tensor(*args) |
| 383 | |
| 384 | |
| 385 | class WeightCacheState(NamedTuple): |
| 386 | """Client-side decode of a daemon's get_state response payload.""" |
| 387 | |
| 388 | entries: dict[str, TensorEntry] |
| 389 | """Model tensors, exported as CUDA IPC handles or shipped by value.""" |
| 390 | aliases: dict[str, str] |
| 391 | """Duplicate (tied) weight names aliased to their canonical entry.""" |
| 392 | attrs: dict[str, bool] |
| 393 | """Python-side flags set by load_weights, e.g. EAGLE ownership flags.""" |
| 394 | |
| 395 | |
| 396 | def send_msg(sock: socket.socket, obj: Any) -> None: |
| 397 | payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) |
| 398 | sock.sendall(_LEN_STRUCT.pack(len(payload))) |
| 399 | sock.sendall(payload) |
| 400 | |
| 401 | |
| 402 | def recv_msg(sock: socket.socket) -> Any: |
| 403 | (length,) = _LEN_STRUCT.unpack(_recv_exact(sock, _LEN_STRUCT.size)) |
| 404 | if length > MAX_MSG_SIZE: |
| 405 | raise ValueError(f"Message size {length} exceeds limit {MAX_MSG_SIZE}") |
| 406 | return pickle.loads(_recv_exact(sock, length)) |
| 407 | |
| 408 | |
| 409 | def _recv_exact(sock: socket.socket, num_bytes: int) -> bytes: |
| 410 | buf = bytearray() |
| 411 | while len(buf) < num_bytes: |
| 412 | chunk = sock.recv(num_bytes - len(buf)) |
| 413 | if not chunk: |
| 414 | raise ConnectionError("Socket closed while receiving message") |
| 415 | buf.extend(chunk) |
| 416 | return bytes(buf) |