| 354 | self.active_update_id: Optional[str] = None |
| 355 | |
| 356 | def install( |
| 357 | self, |
| 358 | manifest: WeightUpdateManifest, |
| 359 | tensors: Mapping[str, torch.Tensor], |
| 360 | ) -> None: |
| 361 | WeightLayout.from_metadata(manifest.metadata).validate_supported() |
| 362 | tensor_map = dict(tensors) |
| 363 | if set(tensor_map) != set(manifest.tensors): |
| 364 | missing = sorted(set(manifest.tensors) - set(tensor_map)) |
| 365 | extra = sorted(set(tensor_map) - set(manifest.tensors)) |
| 366 | raise WeightManifestValidationError( |
| 367 | f"vLLM reload tensor set mismatch: missing={missing}, extra={extra}" |
| 368 | ) |
| 369 | |
| 370 | weights: list[tuple[str, torch.Tensor]] = [] |
| 371 | for name, descriptor in manifest.tensors.items(): |
| 372 | tensor = tensor_map[name] |
| 373 | if tuple(int(dim) for dim in tensor.shape) != descriptor.shape: |
| 374 | raise WeightManifestValidationError( |
| 375 | f"vLLM reload tensor shape mismatch for {name}: " |
| 376 | f"expected {descriptor.shape}, got {tuple(tensor.shape)}" |
| 377 | ) |
| 378 | if str(tensor.dtype) != descriptor.dtype: |
| 379 | raise WeightManifestValidationError( |
| 380 | f"vLLM reload tensor dtype mismatch for {name}: " |
| 381 | f"expected {descriptor.dtype}, got {tensor.dtype}" |
| 382 | ) |
| 383 | |
| 384 | weight = tensor.detach() |
| 385 | if self.target_dtype is not None and weight.dtype != self.target_dtype: |
| 386 | weight = weight.to(dtype=self.target_dtype) |
| 387 | if self.target_device is not None and torch.device(self.target_device) != weight.device: |
| 388 | weight = weight.to(device=self.target_device) |
| 389 | if not weight.is_contiguous(): |
| 390 | weight = weight.contiguous() |
| 391 | weights.append((name, weight)) |
| 392 | |
| 393 | try: |
| 394 | reload_weights = self._resolve_reload_weights() |
| 395 | reload_weights(weights) |
| 396 | if self.synchronize_cuda and torch.cuda.is_available(): |
| 397 | torch.cuda.synchronize() |
| 398 | except Exception: |
| 399 | self.active_update_id = None |
| 400 | raise |
| 401 | |
| 402 | self.active_weight_version = manifest.weight_version |
| 403 | self.active_update_id = manifest.update_id |
| 404 | |
| 405 | def release(self, update_id: str) -> None: |
| 406 | if self.active_update_id == update_id: |