wehub-resource-sync
527 lines
21 KiB
Python
527 lines
21 KiB
Python
import atexit
|
|
import logging
|
|
import requests
|
|
import threading
|
|
import time
|
|
import torch
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import asdict
|
|
from packaging import version
|
|
from pydantic import ValidationError
|
|
from requests import ConnectionError
|
|
from torch import nn
|
|
from typing import List, Optional, Union
|
|
from urllib.parse import urlparse
|
|
|
|
from swift.infer_engine import AdapterRequest, RequestConfig
|
|
from swift.infer_engine.protocol import ChatCompletionResponse, RolloutInferRequest, RolloutOutput
|
|
from swift.metrics import Metric
|
|
from swift.utils import (is_trl_available, is_vllm_ascend_available, is_vllm_available, is_vllm_metax_available,
|
|
synchronize)
|
|
from .utils import (broadcast_tensor_for_vllm_weight_sync, format_host_for_url, is_valid_ipv6_address,
|
|
peft_config_to_dict, resolve_hostname)
|
|
|
|
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
|
|
|
|
if is_vllm_metax_available():
|
|
import vllm_metax.patch
|
|
|
|
if is_trl_available():
|
|
import trl
|
|
trl_verison = version.parse(trl.__version__)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class VLLMInferClient:
|
|
"""Inference-only vLLM client. Posts to /infer/ endpoint.
|
|
No weight synchronization. Used for GKD teacher server etc.
|
|
"""
|
|
|
|
def __init__(self,
|
|
base_urls: Optional[List[str]] = None,
|
|
hosts: List[str] = ['0.0.0.0'],
|
|
server_ports: List[int] = [8000],
|
|
connection_timeout: float = 240.0):
|
|
if not is_vllm_available():
|
|
raise ImportError('vLLM is not installed. Please install it with `pip install vllm`.')
|
|
|
|
if base_urls is not None:
|
|
self.base_urls = []
|
|
self.hosts = []
|
|
for url in base_urls:
|
|
parsed_url = urlparse(url)
|
|
# Use resolve_hostname instead of gethostbyname for IPv6 support
|
|
host = resolve_hostname(parsed_url.hostname)
|
|
scheme = parsed_url.scheme or 'http'
|
|
base_url_i = f'{scheme}://{parsed_url.netloc}{parsed_url.path}'
|
|
self.base_urls.append(base_url_i)
|
|
self.hosts.append(host)
|
|
else:
|
|
if len(hosts) != len(server_ports):
|
|
raise ValueError('host and server_port must have same length when lists are provided')
|
|
# Format IPv6 addresses correctly in URLs (wrap with brackets)
|
|
self.base_urls = [f'http://{format_host_for_url(h)}:{p}' for h, p in zip(hosts, server_ports)]
|
|
self.hosts = hosts
|
|
|
|
self.num_servers = len(self.base_urls)
|
|
self.sessions = [requests.Session() for _ in range(self.num_servers)]
|
|
self.check_server(connection_timeout)
|
|
|
|
def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0):
|
|
server_status = [False] * self.num_servers
|
|
|
|
def check_single_server(i):
|
|
start_time = time.time()
|
|
url = f'{self.base_urls[i]}/health/'
|
|
while True:
|
|
try:
|
|
response = requests.get(url, timeout=retry_interval)
|
|
if response.status_code == 200:
|
|
server_status[i] = True
|
|
return
|
|
except Exception:
|
|
pass
|
|
|
|
elapsed = time.time() - start_time
|
|
if elapsed >= total_timeout:
|
|
return
|
|
|
|
time.sleep(retry_interval)
|
|
|
|
threads = []
|
|
for i in range(self.num_servers):
|
|
t = threading.Thread(target=check_single_server, args=(i, ))
|
|
t.daemon = True
|
|
t.start()
|
|
threads.append(t)
|
|
|
|
for t in threads:
|
|
t.join(total_timeout)
|
|
|
|
if not all(server_status):
|
|
failed_servers = [self.base_urls[i] for i, status in enumerate(server_status) if not status]
|
|
raise ConnectionError(f'Servers not reachable after {total_timeout}s: {failed_servers}')
|
|
|
|
def infer(
|
|
self,
|
|
infer_requests: List[RolloutInferRequest],
|
|
request_config: Optional[RequestConfig] = None,
|
|
metrics: Optional[List[Metric]] = None,
|
|
*,
|
|
use_tqdm: Optional[bool] = None,
|
|
adapter_request: Optional[AdapterRequest] = None,
|
|
):
|
|
n = len(infer_requests)
|
|
chunk_size = (n + self.num_servers - 1) // self.num_servers
|
|
chunks = [infer_requests[i:i + chunk_size] for i in range(0, n, chunk_size)]
|
|
chunks += [[]] * (self.num_servers - len(chunks))
|
|
|
|
results = [None] * self.num_servers
|
|
errors = [None] * self.num_servers
|
|
if isinstance(request_config, RequestConfig):
|
|
request_config = asdict(request_config)
|
|
|
|
def process_chunk(i, chunk):
|
|
try:
|
|
if len(chunk) > 0 and isinstance(chunk[0], RolloutInferRequest):
|
|
chunk = [asdict(req) for req in chunk]
|
|
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/infer/',
|
|
json={
|
|
'infer_requests': chunk,
|
|
'request_config': request_config,
|
|
'metrics': metrics,
|
|
'use_tqdm': use_tqdm,
|
|
'adapter_request': adapter_request,
|
|
},
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
errors[i] = Exception(f'Server {i} failed: {response.status_code}, {response.text}')
|
|
return
|
|
|
|
resp_data = response.json()
|
|
parsed: List[Union[RolloutOutput, ChatCompletionResponse]] = []
|
|
|
|
for item in resp_data:
|
|
try:
|
|
parsed.append(RolloutOutput.model_validate(item))
|
|
except ValidationError:
|
|
parsed.append(ChatCompletionResponse(**item))
|
|
results[i] = parsed
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(process_chunk, i, chunk) for i, chunk in enumerate(chunks)]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors: {all_errors}')
|
|
|
|
return [res for server_results in results for res in server_results]
|
|
|
|
|
|
class VLLMClient(VLLMInferClient):
|
|
|
|
def __init__(self,
|
|
base_urls: Optional[List[str]] = None,
|
|
hosts: List[str] = ['0.0.0.0'],
|
|
server_ports: List[int] = [8000],
|
|
group_ports: Optional[Union[int, List[int]]] = None,
|
|
connection_timeout: float = 240.0):
|
|
super().__init__(
|
|
base_urls=base_urls, hosts=hosts, server_ports=server_ports, connection_timeout=connection_timeout)
|
|
|
|
if group_ports is None:
|
|
group_ports = [51216 + i for i in range(self.num_servers)]
|
|
|
|
if isinstance(group_ports, int):
|
|
self.group_ports = [group_ports + i for i in range(self.num_servers)]
|
|
elif isinstance(group_ports, list) and len(group_ports) == self.num_servers:
|
|
self.group_ports = group_ports
|
|
else:
|
|
raise ValueError('group_port must be int or list of length num_servers')
|
|
|
|
self.pynccl_comms = []
|
|
|
|
def init_communicator(self, device: Union[int, str] = 0):
|
|
self.pynccl_comms = []
|
|
for i in range(self.num_servers):
|
|
response = self.sessions[i].get(f'{self.base_urls[i]}/get_world_size/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} failed: {response.text}')
|
|
vllm_world_size = response.json()['world_size']
|
|
|
|
world_size = vllm_world_size + 1
|
|
rank = vllm_world_size
|
|
# Use '::' for IPv6 hosts, '0.0.0.0' for IPv4 hosts
|
|
bind_host = '::' if is_valid_ipv6_address(self.hosts[i]) else '0.0.0.0'
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/init_communicator/',
|
|
json={
|
|
'host': bind_host,
|
|
'port': self.group_ports[i],
|
|
'world_size': world_size,
|
|
})
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} init failed: {response.text}')
|
|
|
|
time.sleep(0.1)
|
|
|
|
pg = StatelessProcessGroup.create(
|
|
host=self.hosts[i], port=self.group_ports[i], rank=rank, world_size=world_size)
|
|
if is_vllm_ascend_available():
|
|
import torch_npu
|
|
torch_npu.npu.set_device(device)
|
|
comm = PyNcclCommunicator(pg, device=device)
|
|
self.pynccl_comms.append(comm)
|
|
|
|
atexit.register(self.close_communicator)
|
|
|
|
def update_named_param(self, name: str, weights: torch.Tensor):
|
|
dtype = str(weights.dtype)
|
|
shape = tuple(weights.shape)
|
|
|
|
errors = [None] * self.num_servers
|
|
|
|
def _update_single_server(i):
|
|
try:
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/update_named_param/',
|
|
json={
|
|
'name': name,
|
|
'dtype': dtype,
|
|
'shape': shape
|
|
},
|
|
)
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} update failed: {response.text}')
|
|
|
|
synchronize()
|
|
broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], weights, src=self.pynccl_comms[i].rank)
|
|
synchronize()
|
|
self.pynccl_comms[i].group.barrier()
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors: {all_errors}')
|
|
|
|
def update_adapter_flattened_param(self, peft_config, metadatas, flattened_tensor):
|
|
"""
|
|
Adds a LoRA adapter to the model on all servers using flattened tensor.
|
|
|
|
Args:
|
|
peft_config: PEFT configuration for LoRA adapter.
|
|
metadatas: List of FlattenedTensorMetadata objects.
|
|
flattened_tensor: The flattened tensor containing all adapter parameters.
|
|
"""
|
|
errors = [None] * self.num_servers
|
|
peft_config = peft_config_to_dict(peft_config)
|
|
metadatas = [m.model_dump() if hasattr(m, 'model_dump') else m.dict() for m in metadatas]
|
|
|
|
def _update_single_server(i):
|
|
try:
|
|
data = {
|
|
'peft_config': {
|
|
**peft_config
|
|
},
|
|
'metadatas': metadatas,
|
|
}
|
|
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/update_adapter_flattened_param/',
|
|
json=data,
|
|
)
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} update adapter failed: {response.text}')
|
|
|
|
synchronize()
|
|
broadcast_tensor_for_vllm_weight_sync(
|
|
self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank)
|
|
synchronize()
|
|
self.pynccl_comms[i].group.barrier()
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors: {all_errors}')
|
|
|
|
def update_adapter_param(self, peft_config, lora_params):
|
|
"""
|
|
Adds a LoRA adapter to the model on all servers without flattening.
|
|
Sends each tensor individually.
|
|
|
|
Args:
|
|
peft_config: PEFT configuration for LoRA adapter.
|
|
lora_params: OrderedDict of (name, tensor) pairs for LoRA parameters.
|
|
"""
|
|
errors = [None] * self.num_servers
|
|
peft_config = peft_config_to_dict(peft_config)
|
|
|
|
# Build metadata for each tensor
|
|
lora_tensors_metadata = []
|
|
for name, param in lora_params.items():
|
|
metadata = {
|
|
'name': name,
|
|
'dtype': str(param.dtype),
|
|
'shape': tuple(param.shape),
|
|
'start_idx': 0, # Not used in non-flattened mode
|
|
'end_idx': param.numel(), # Not used in non-flattened mode
|
|
'numel': param.numel(),
|
|
}
|
|
lora_tensors_metadata.append(metadata)
|
|
|
|
def _update_single_server(i):
|
|
try:
|
|
data = {
|
|
'peft_config': {
|
|
**peft_config
|
|
},
|
|
'lora_tensors_metadata': lora_tensors_metadata,
|
|
}
|
|
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/update_adapter_param/',
|
|
json=data,
|
|
)
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} update adapter failed: {response.text}')
|
|
|
|
# Broadcast each tensor individually
|
|
synchronize()
|
|
for name, param in lora_params.items():
|
|
broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], param, src=self.pynccl_comms[i].rank)
|
|
synchronize()
|
|
self.pynccl_comms[i].group.barrier()
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors: {all_errors}')
|
|
|
|
def update_flattened_params(self, metadatas, flattened_tensor):
|
|
"""
|
|
Updates model parameters using flattened tensor data.
|
|
|
|
Args:
|
|
metadatas: List of FlattenedTensorMetadata objects
|
|
flattened_tensor: The flattened tensor containing all parameters
|
|
"""
|
|
errors = [None] * self.num_servers
|
|
metadatas = [m.model_dump() if hasattr(m, 'model_dump') else m.dict() for m in metadatas]
|
|
|
|
def _update_single_server(i):
|
|
try:
|
|
data = {
|
|
'metadatas': metadatas,
|
|
}
|
|
|
|
response = self.sessions[i].post(
|
|
f'{self.base_urls[i]}/update_flattened_params/',
|
|
json=data,
|
|
)
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} update flattened params failed: {response.text}')
|
|
|
|
synchronize()
|
|
broadcast_tensor_for_vllm_weight_sync(
|
|
self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank)
|
|
synchronize()
|
|
self.pynccl_comms[i].group.barrier()
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors: {all_errors}')
|
|
|
|
def process_weights_after_loading(self):
|
|
"""Trigger process_weights_after_loading on all vLLM workers.
|
|
|
|
Must be called **once** after ALL weight buckets have been
|
|
sent via ``update_flattened_params``. This mirrors the
|
|
pattern used by verl and ROLL.
|
|
"""
|
|
errors = [None] * self.num_servers
|
|
|
|
def _process_single_server(i):
|
|
try:
|
|
response = self.sessions[i].post(f'{self.base_urls[i]}/process_weights_after_loading/', )
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} process_weights_after_loading failed: {response.text}')
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_process_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors on process_weights_after_loading: {all_errors}')
|
|
|
|
def update_model_params(self, model: nn.Module):
|
|
for name, param in model.named_parameters():
|
|
self.update_named_param(name, param.data)
|
|
|
|
def reset_prefix_cache(self):
|
|
errors = [None] * self.num_servers
|
|
|
|
def _reset_single_server(i):
|
|
try:
|
|
response = self.sessions[i].post(f'{self.base_urls[i]}/reset_prefix_cache/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} reset failed: {response.text}')
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors on reset_prefix_cache: {all_errors}')
|
|
|
|
def reset_encoder_cache(self):
|
|
errors = [None] * self.num_servers
|
|
|
|
def _reset_single_server(i):
|
|
try:
|
|
response = self.sessions[i].post(f'{self.base_urls[i]}/reset_encoder_cache/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} reset failed: {response.text}')
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors on reset_encoder_cache: {all_errors}')
|
|
|
|
def reset_mm_cache(self):
|
|
errors = [None] * self.num_servers
|
|
|
|
def _reset_single_server(i):
|
|
try:
|
|
response = self.sessions[i].post(f'{self.base_urls[i]}/reset_mm_cache/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Server {i} reset failed: {response.text}')
|
|
except Exception as e:
|
|
errors[i] = e
|
|
|
|
with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
|
|
futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
|
|
for future in futures:
|
|
future.result()
|
|
all_errors = [e for e in errors if e is not None]
|
|
if all_errors:
|
|
raise RuntimeError(f'Multiple errors on reset_mm_cache: {all_errors}')
|
|
|
|
def get_engine_type(self):
|
|
# assume that all server has same engine type
|
|
response = self.sessions[0].post(f'{self.base_urls[0]}/get_engine_type/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Engine type request failed: {response.text}')
|
|
|
|
result = response.json()
|
|
self.use_async_engine = result['engine_type'] == 'AsyncLLMEngine'
|
|
self.enable_multi_turn = result.get('enable_multi_turn', False)
|
|
self.use_gym_env = result.get('use_gym_env', False)
|
|
self.enable_lora = result.get('enable_lora', False)
|
|
return result
|
|
|
|
def get_model_state_keys(self):
|
|
"""Fetch runtime vLLM model parameter names from server."""
|
|
response = self.sessions[0].get(f'{self.base_urls[0]}/get_model_state_keys/')
|
|
if response.status_code != 200:
|
|
raise Exception(f'Get model state keys failed: {response.text}')
|
|
data = response.json()
|
|
return data.get('keys', [])
|
|
|
|
def close_communicator(self):
|
|
for i in range(self.num_servers):
|
|
try:
|
|
response = self.sessions[i].post(f'{self.base_urls[i]}/close_communicator/')
|
|
if response.status_code != 200:
|
|
logger.warning(f'Server {i} close failed: {response.text}')
|
|
except Exception as e:
|
|
logger.warning(f'Error closing server {i} communicator: {str(e)}')
|