| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Built-in `VLLMWeightSyncClient` implementations. |
| 4 | |
| 5 | These adapt the inference engine's weight-sync control plane to concrete |
| 6 | transports. A `TrainerWeightTransferEngine` takes one of these (or any object |
| 7 | with the same four methods — the protocol is structural) and drives the full |
| 8 | handshake through it. |
| 9 | |
| 10 | Imports of `ray` / `requests` are deferred to call time so this module is |
| 11 | importable without those packages installed. |
| 12 | """ |
| 13 | |
| 14 | from typing import TYPE_CHECKING, Any |
| 15 | |
| 16 | from vllm.distributed.weight_transfer.base import ( |
| 17 | WeightTransferInitRequest, |
| 18 | WeightTransferUpdatePayload, |
| 19 | WeightTransferUpdateRequest, |
| 20 | ) |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from ray.actor import ActorHandle |
| 24 | |
| 25 | |
| 26 | def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]: |
| 27 | """Make an update_info dict JSON-serializable for HTTP transport. |
| 28 | |
| 29 | CUDA IPC handles (`ipc_handles`) are tuples of non-JSON-native objects, so |
| 30 | over HTTP they are pickled+base64-encoded into `ipc_handles_pickled` (which |
| 31 | the worker auto-deserializes when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`). |
| 32 | Other backends (NCCL) carry only JSON-native metadata and pass through |
| 33 | unchanged. Mirrors the old IPC `_do_send` HTTP branch. |
| 34 | """ |
| 35 | ipc_handles = update_info.get("ipc_handles") |
| 36 | if ipc_handles is None: |
| 37 | return update_info |
| 38 | |
| 39 | import pickle |
| 40 | |
| 41 | import pybase64 as base64 |
| 42 | |
| 43 | out = {k: v for k, v in update_info.items() if k != "ipc_handles"} |
| 44 | out["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(ipc_handles)).decode( |
| 45 | "utf-8" |
| 46 | ) |
| 47 | return out |
| 48 | |
| 49 | |
| 50 | def _json_safe_update_payload( |
| 51 | update_info: WeightTransferUpdatePayload, |
| 52 | ) -> WeightTransferUpdatePayload: |
| 53 | if isinstance(update_info, list): |
| 54 | return [_json_safe_update_info(info) for info in update_info] |
| 55 | return _json_safe_update_info(update_info) |
| 56 | |
| 57 | |
| 58 | class HTTPVLLMWeightSyncClient: |
| 59 | """Talks to a vLLM server over the RLHF HTTP routes. |
| 60 | |
| 61 | Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`: |
| 62 | `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, |
| 63 | `/finish_weight_update`. |
| 64 | """ |
| 65 | |
| 66 | def __init__(self, base_url: str, timeout: float = 300) -> None: |
| 67 | self.base_url = base_url.rstrip("/") |
| 68 | self.timeout = timeout |
| 69 | |
| 70 | def _post(self, path: str, json: dict[str, Any] | None = None) -> None: |
| 71 | import requests |
| 72 | |
| 73 | response = requests.post( |
| 74 | f"{self.base_url}/{path}", json=json, timeout=self.timeout |
| 75 | ) |
| 76 | response.raise_for_status() |
| 77 | |
| 78 | def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: |
| 79 | self._post("init_weight_transfer_engine", {"init_info": init_info}) |
| 80 | |
| 81 | def start_weight_update(self) -> None: |
| 82 | self._post("start_weight_update") |
| 83 | |
| 84 | def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: |
| 85 | self._post( |
| 86 | "update_weights", {"update_info": _json_safe_update_payload(update_info)} |
| 87 | ) |
| 88 | |
| 89 | def finish_weight_update(self, weight_version: str | None = None) -> None: |
| 90 | json = ( |
| 91 | {"weight_version": weight_version} if weight_version is not None else None |
| 92 | ) |
| 93 | self._post("finish_weight_update", json) |
| 94 | |
| 95 | |
| 96 | class RayVLLMWeightSyncClient: |
| 97 | """Talks to one or more vLLM `AsyncLLM`/`LLM` Ray actors. |
| 98 | |
| 99 | Each call fans out to every handle and blocks on all of them, so a |
| 100 | multi-actor (e.g. multi-DP) deployment is driven as one unit. |
| 101 | """ |
| 102 | |
| 103 | def __init__(self, handle: "ActorHandle | list[ActorHandle]") -> None: |
| 104 | self.handles = handle if isinstance(handle, list) else [handle] |
| 105 | |
| 106 | def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: |
| 107 | import ray |
| 108 | |
| 109 | request = WeightTransferInitRequest(init_info=init_info) |
| 110 | ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles]) |
| 111 | |
| 112 | def start_weight_update(self) -> None: |
| 113 | import ray |
| 114 | |
| 115 | ray.get([h.start_weight_update.remote() for h in self.handles]) |
| 116 | |
| 117 | def update_weights(self, update_info: WeightTransferUpdatePayload) -> None: |
| 118 | import ray |
| 119 | |
| 120 | request = WeightTransferUpdateRequest(update_info=update_info) |
| 121 | ray.get([h.update_weights.remote(request) for h in self.handles]) |
| 122 | |
| 123 | def finish_weight_update(self, weight_version: str | None = None) -> None: |
| 124 | import ray |
| 125 | |
| 126 | ray.get([h.finish_weight_update.remote() for h in self.handles]) |
| 127 | if weight_version is not None: |
| 128 | ray.get( |
| 129 | [h.update_weight_version.remote(weight_version) for h in self.handles] |
| 130 | ) |