MCPcopy Create free account
hub / github.com/RL-Align/RL-Kernel / VLLMCheckpointWeightReloadAdapter

Class VLLMCheckpointWeightReloadAdapter

rl_engine/executors/bridge.py:436–540  ·  view source on GitHub ↗

Install a manifest through vLLM's checkpoint-path reload utility path. This path works with vLLM's default EngineCore multiprocessing because the worker process reloads weights from `weights_path` itself; no CUDA tensor is serialized across the RPC boundary. It is a production-alig

Source from the content-addressed store, hash-verified

434
435
436class VLLMCheckpointWeightReloadAdapter:
437 """
438 Install a manifest through vLLM's checkpoint-path reload utility path.
439
440 This path works with vLLM's default EngineCore multiprocessing because the
441 worker process reloads weights from `weights_path` itself; no CUDA tensor is
442 serialized across the RPC boundary. It is a production-aligned hot reload
443 fallback for environments where CUDA IPC/NCCL transport is unavailable, but
444 it is not a zero-copy transport.
445 """
446
447 def __init__(
448 self,
449 engine: Any,
450 *,
451 weights_path: Optional[str] = None,
452 weights_path_resolver: Optional[Any] = None,
453 metadata_key: str = "vllm_weights_path",
454 is_checkpoint_format: bool = True,
455 synchronize_cuda: bool = True,
456 ):
457 self.engine = engine
458 self.weights_path = weights_path
459 self.weights_path_resolver = weights_path_resolver
460 self.metadata_key = metadata_key
461 self.is_checkpoint_format = bool(is_checkpoint_format)
462 self.synchronize_cuda = bool(synchronize_cuda)
463 self.active_weight_version: Optional[int] = None
464 self.active_update_id: Optional[str] = None
465
466 def install(
467 self,
468 manifest: WeightUpdateManifest,
469 tensors: Mapping[str, torch.Tensor],
470 ) -> None:
471 WeightLayout.from_metadata(manifest.metadata).validate_supported()
472 tensor_map = dict(tensors)
473 if set(tensor_map) != set(manifest.tensors):
474 missing = sorted(set(manifest.tensors) - set(tensor_map))
475 extra = sorted(set(tensor_map) - set(manifest.tensors))
476 raise WeightManifestValidationError(
477 f"vLLM checkpoint reload tensor set mismatch: missing={missing}, extra={extra}"
478 )
479
480 weights_path = self._resolve_weights_path(manifest, tensor_map)
481 try:
482 reload_weights = self._resolve_reload_weights()
483 reload_weights(weights_path)
484 if self.synchronize_cuda and torch.cuda.is_available():
485 torch.cuda.synchronize()
486 except Exception:
487 self.active_update_id = None
488 raise
489
490 self.active_weight_version = manifest.weight_version
491 self.active_update_id = manifest.update_id
492
493 def release(self, update_id: str) -> None:

Calls

no outgoing calls