File size: 8,984 Bytes
7feac49 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
# Copyright (c) Alibaba, Inc. and its affiliates.
# Code partially sourced from Hugging Face TRL
import atexit
import logging
import time
from typing import List, Optional
import requests
import torch
from dacite import from_dict
from requests import ConnectionError
from torch import nn
from swift.llm import AdapterRequest, InferRequest, Template
from swift.llm.infer.protocol import ChatCompletionResponse, RequestConfig
from swift.plugin import Metric
from swift.utils import is_vllm_ascend_available, is_vllm_available
if is_vllm_available():
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
if is_vllm_ascend_available():
from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator # noqa
logger = logging.getLogger(__name__)
class VLLMClient:
"""
A client class to interact with a vLLM server.
This class provides methods to infer completions, initialize and manage weight update groups, and update model
weights in a distributed setting. Before using it, start the vLLM server with `trl vllm-serve`.
Args:
host (`str`, *optional*, defaults to `"0.0.0.0"`):
IP address of the vLLM server.
server_port (`int`, *optional*, defaults to `8000`):
Port number of the vLLM server.
group_port (`int`, *optional*, defaults to `51216`):
Port number for the weight update group.
connection_timeout (`float`, *optional*, defaults to `0.0`):
Total timeout duration in seconds to wait for the server to be up. If the server is not up after the
timeout, a `ConnectionError` is raised.
"""
def __init__(self,
host: str = '0.0.0.0',
server_port: int = 8000,
group_port: int = 51216,
connection_timeout: float = 0.0):
if not is_vllm_available():
raise ImportError('vLLM is not installed. Please install it with `pip install vllm`.')
self.session = requests.Session()
self.host = host
self.server_port = server_port
self.group_port = group_port
self.check_server(connection_timeout) # check server and fail after timeout
def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0):
"""
Check server availability with retries on failure, within a total timeout duration. If the server is not up
after the total timeout duration, raise a `ConnectionError`.
Args:
retry_interval (`float`, *optional*, defaults to `2.0`):
Interval in seconds between retries.
total_timeout (`float`, *optional*, defaults to `0.0`):
Total timeout duration in seconds.
"""
url = f'http://{self.host}:{self.server_port}/health/'
start_time = time.time() # Record the start time
while True:
try:
response = requests.get(url)
except requests.exceptions.RequestException as exc:
# Check if the total timeout duration has passed
elapsed_time = time.time() - start_time
if elapsed_time >= total_timeout:
raise ConnectionError(
f"The vLLM server can't be reached at {self.host}:{self.server_port} after {total_timeout} "
'seconds. Make sure the server is running by running `swift deploy`.') from exc
else:
if response.status_code == 200:
logger.info('Server is up!')
return None
# Retry logic: wait before trying again
logger.info(f'Server is not up yet. Retrying in {retry_interval} seconds...')
time.sleep(retry_interval)
def infer(
self,
infer_requests: List[InferRequest],
request_config: Optional[RequestConfig] = None,
metrics: Optional[List[Metric]] = None,
*,
template: Optional[Template] = None,
use_tqdm: Optional[bool] = None,
adapter_request: Optional[AdapterRequest] = None,
):
url = f'http://{self.host}:{self.server_port}/infer/'
response = self.session.post(
url,
json={
'infer_requests': infer_requests,
'request_config': request_config,
'metrics': metrics,
'template': template,
'use_tqdm': use_tqdm,
'adapter_request': adapter_request,
},
)
if response.status_code == 200:
return [from_dict(data_class=ChatCompletionResponse, data=resp) for resp in response.json()]
else:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
def init_communicator(self):
"""
Initializes the weight update group in a distributed setup for model synchronization.
"""
# Get the tensor parallel size from the server
url = f'http://{self.host}:{self.server_port}/get_world_size/'
response = requests.get(url)
if response.status_code == 200:
vllm_world_size = response.json()['world_size']
else:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
world_size = vllm_world_size + 1 # add the client to the world
self.rank = vllm_world_size # the client's rank is the last process
# Initialize weight update group
url = f'http://{self.host}:{self.server_port}/init_communicator/'
# In the server side, the host is set to 0.0.0.0
response = self.session.post(url, json={'host': '0.0.0.0', 'port': self.group_port, 'world_size': world_size})
if response.status_code != 200:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
# Brief delay to allow server initialization. While not strictly required (client socket will retry on
# connection failure), this prevents log warnings like:
# [W416 23:24:57.460001114 socket.cpp:204] [c10d] The hostname of the client socket cannot be retrieved. err=-3
time.sleep(0.1)
# Set up the communication group for weight broadcasting
pg = StatelessProcessGroup.create(host=self.host, port=self.group_port, rank=self.rank, world_size=world_size)
self.pynccl_comm = PyNcclCommunicator(pg, device=0)
# When the client object is deleted, close the weight update group
atexit.register(self.close_communicator)
def update_named_param(self, name: str, weights: torch.Tensor):
"""
Updates a specific named parameter in the model and broadcasts it to other processes.
Args:
name (`str`):
Name of the layer whose weights are being updated.
weights (`torch.Tensor`):
Tensor containing the updated weights.
"""
dtype, shape = str(weights.dtype), tuple(weights.shape)
url = f'http://{self.host}:{self.server_port}/update_named_param/'
response = self.session.post(url, json={'name': name, 'dtype': dtype, 'shape': shape})
if response.status_code != 200:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
# Broadcast the weights to the other processes
self.pynccl_comm.broadcast(weights, src=self.rank)
self.pynccl_comm.group.barrier()
def update_model_params(self, model: nn.Module):
"""
Updates all parameters of the given model by calling `update_named_param` for each parameter in the model.
Args:
model (`nn.Module`):
Model whose parameters (weights/biases) are to be updated.
"""
for name, param in model.named_parameters():
# Update each parameter individually
self.update_named_param(name, param.data)
def reset_prefix_cache(self):
"""
Resets the prefix cache for the model.
"""
url = f'http://{self.host}:{self.server_port}/reset_prefix_cache/'
response = self.session.post(url)
if response.status_code != 200:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
def close_communicator(self):
"""
Closes the weight update group and cleans up the communication group.
"""
url = f'http://{self.host}:{self.server_port}/close_communicator/'
try:
response = self.session.post(url)
except ConnectionError:
# The server might be already down, so we don't need to close the communicator
pass
else:
if response.status_code != 200:
raise Exception(f'Request failed: {response.status_code}, {response.text}')
|