423 lines
14 KiB
Python
423 lines
14 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import json
|
|
import socket
|
|
import threading
|
|
from contextlib import suppress
|
|
|
|
import regex as re
|
|
import torch
|
|
from vllm.logger import logger
|
|
|
|
from ..executor.elastic_load import P2PSend
|
|
from ..utils import find_free_port
|
|
|
|
|
|
class ElasticClient:
|
|
"""
|
|
Class for handling the client-side logic of Netloader of models.
|
|
"""
|
|
|
|
def __init__(
|
|
self, sources: list[str], device_id: int, model_path: str, tp: int, pp: int, group_name: str = "netloader"
|
|
):
|
|
"""
|
|
Initializes the ElasticClient instance.
|
|
|
|
Parameters:
|
|
- sources: List of source addresses in the format IP:port.
|
|
- device_id: The ID of the current device.
|
|
- model_path: The path to the model.
|
|
- tp: Tensor parallel size.
|
|
- pp: Pipeline parallel size.
|
|
- group_name: Name of the HCCL process group.
|
|
"""
|
|
self.sources = sources
|
|
self.device_id = device_id
|
|
self.model_path = model_path
|
|
self.tp = tp
|
|
self.pp = pp
|
|
self.group_name = group_name
|
|
|
|
self.s: socket.socket | None = None
|
|
self.ack: tuple[str, int] | None = None
|
|
self.server_addr: str | None = None
|
|
self.server_port: int | None = None
|
|
|
|
for source in self.sources:
|
|
try:
|
|
ip, port_str = source.split(":")
|
|
port = int(port_str)
|
|
except Exception as e:
|
|
logger.info("IP format error: %s, detail: %s", source, e)
|
|
continue
|
|
|
|
self.server_addr = ip
|
|
self.server_port = port
|
|
|
|
try:
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
logger.info("Start connection to server: %s:%s", self.server_addr, self.server_port)
|
|
sock.connect((self.server_addr, self.server_port))
|
|
logger.info("Finish connection to server: %s:%s", self.server_addr, self.server_port)
|
|
sock.settimeout(60)
|
|
|
|
self.s = sock
|
|
self.ack = self.register(device_id, model_path, tp, pp)
|
|
break
|
|
except Exception as e:
|
|
logger.error("Connect to %s fails, detail: %s", source, e)
|
|
if sock is not None:
|
|
with suppress(Exception):
|
|
sock.close()
|
|
self.s = None
|
|
self.ack = None
|
|
self.server_addr = None
|
|
self.server_port = None
|
|
|
|
if self.s is None:
|
|
sources_str = ", ".join(self.sources[:2])
|
|
if len(self.sources) > 2:
|
|
sources_str += f", ... (total {len(self.sources)})"
|
|
logger.error(
|
|
"All sources exhausted, no connection established for device_id=%s, model_path=%s, sources=[%s]",
|
|
device_id,
|
|
model_path,
|
|
sources_str,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Closes the socket connection.
|
|
"""
|
|
if self.s is not None:
|
|
try:
|
|
self.s.close()
|
|
except Exception as e:
|
|
logger.error("Error closing socket: %s", e)
|
|
finally:
|
|
self.s = None
|
|
|
|
def __enter__(self) -> "ElasticClient":
|
|
"""
|
|
Context manager enter method.
|
|
"""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
"""
|
|
Context manager exit method.
|
|
"""
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
"""
|
|
Destructor method to ensure socket is closed.
|
|
"""
|
|
with suppress(Exception):
|
|
self.close()
|
|
|
|
def send_str(self, data_str: str) -> None:
|
|
"""
|
|
Sends a string over the socket connection.
|
|
|
|
Parameters:
|
|
- data_str: The string to be sent.
|
|
"""
|
|
if self.s is None:
|
|
raise RuntimeError("Socket was not created correctly.")
|
|
self.s.send(data_str.encode("utf-8"))
|
|
|
|
def recv_str(self, buffer_size: int = 1024) -> str:
|
|
"""
|
|
Receives a string over the socket connection.
|
|
|
|
Parameters:
|
|
- buffer_size: The size of the buffer for receiving data.
|
|
|
|
Returns:
|
|
- The received string.
|
|
"""
|
|
if self.s is None:
|
|
raise RuntimeError("Socket was not created correctly.")
|
|
data_str = self.s.recv(buffer_size).decode("utf-8")
|
|
return data_str
|
|
|
|
def register(self, device_id: int, model_path: str, tp: int, pp: int) -> tuple[str, int]:
|
|
"""
|
|
Registers the client with the server.
|
|
|
|
Parameters:
|
|
- device_id: The ID of the current device.
|
|
- model_path: The path to the model.
|
|
- tp: Tensor parallel size.
|
|
- pp: Pipeline parallel size.
|
|
|
|
Returns:
|
|
- A tuple containing the communication name and port.
|
|
"""
|
|
free_port = find_free_port()
|
|
data = {
|
|
"label": "JOIN",
|
|
"content": {
|
|
"device_id": device_id,
|
|
"model_path": model_path,
|
|
"tp": tp,
|
|
"pp": pp,
|
|
"port": free_port,
|
|
"group_name": self.group_name,
|
|
},
|
|
}
|
|
|
|
try:
|
|
self.send_str(json.dumps(data))
|
|
except Exception as e:
|
|
raise RuntimeError(f"Send data {data} to server fails, detail: {e}")
|
|
|
|
try:
|
|
ack_str = self.recv_str()
|
|
except Exception as e:
|
|
raise RuntimeError(f"Receive data from server fails, detail: {e}")
|
|
|
|
try:
|
|
ack = json.loads(ack_str)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Receive data {ack_str} cannot be converted to JSON format, detail: {e}")
|
|
|
|
logger.info("Receive ack: %s", ack)
|
|
|
|
if (
|
|
"label" in ack
|
|
and ack["label"] == "JOIN_ACK"
|
|
and "content" in ack
|
|
and ack["content"] is not None
|
|
and "name" in ack["content"]
|
|
):
|
|
return (ack["content"]["name"], free_port)
|
|
elif "label" in ack and ack["label"] == "JOIN_NACK" and "content" in ack:
|
|
raise RuntimeError(f"Receive nack from server, reason: {ack['content']}")
|
|
else:
|
|
raise RuntimeError(f"Receive ack {ack} from server does not contain required fields")
|
|
|
|
|
|
class ElasticServer:
|
|
"""
|
|
Class for handling the server-side logic of Netloader of models.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
addr: str,
|
|
port: int,
|
|
model,
|
|
device_id: int,
|
|
model_path: str,
|
|
tp: int,
|
|
pp: int,
|
|
int8_cache: str,
|
|
int8_cache_name: list[str] | None,
|
|
group_name: str = "netloader",
|
|
):
|
|
"""
|
|
Initializes the ElasticServer instance.
|
|
|
|
Parameters:
|
|
- addr: The IP address to listen on.
|
|
- port: The port number to listen on.
|
|
- model: The model to be served.
|
|
- device_id: The ID of the current device (i.e. global rank).
|
|
- model_path: The path to the model.
|
|
- tp: Tensor parallel size.
|
|
- pp: Pipeline parallel size.
|
|
- int8_cache: The type of caching for int8 parameters (HBM, DRAM, or no).
|
|
- int8_cache_name: List of parameter names to be cached.
|
|
- group_name: Name of the HCCL process group.
|
|
"""
|
|
self.addr = addr
|
|
self.port = port
|
|
self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
self.s.bind((self.addr, self.port))
|
|
self.s.listen(256)
|
|
|
|
self.model = model
|
|
self.device_id = device_id
|
|
self.model_path = model_path
|
|
self.tp = tp
|
|
self.pp = pp
|
|
self.group_name = group_name
|
|
|
|
self.original_int8 = {}
|
|
int8_pattern = "|".join(map(re.escape, int8_cache_name)) if int8_cache_name is not None else "(?:)"
|
|
for name, param in self.model.named_parameters():
|
|
if param.dtype == torch.int8:
|
|
if int8_cache == "hbm":
|
|
if int8_cache_name is None or (
|
|
int8_cache_name is not None and re.search(int8_pattern, name) is not None
|
|
):
|
|
try:
|
|
self.original_int8[name] = param.data.clone().detach()
|
|
except RuntimeError as e:
|
|
logger.error("Failed to cache int8 tensor %s to HBM, change to DRAM, due to %s", name, e)
|
|
self.original_int8[name] = param.data.cpu()
|
|
|
|
elif int8_cache == "dram":
|
|
if int8_cache_name is None or (
|
|
int8_cache_name is not None and re.search(int8_pattern, name) is not None
|
|
):
|
|
self.original_int8[name] = param.data.cpu()
|
|
elif int8_cache == "no":
|
|
pass
|
|
else:
|
|
logger.warning(
|
|
"int8_cache should be selected in [HBM, DRAM], but got %s, change to no cache", int8_cache
|
|
)
|
|
|
|
logger.info(
|
|
"Server %s:%s starts, device id: %s, model path: %s, tp: %s, pp: %s, int8 params %s are saved to %s",
|
|
self.addr,
|
|
self.port,
|
|
self.device_id,
|
|
self.model_path,
|
|
self.tp,
|
|
self.pp,
|
|
list(self.original_int8),
|
|
int8_cache,
|
|
)
|
|
|
|
def __del__(self):
|
|
"""
|
|
Destructor method to ensure socket is closed.
|
|
"""
|
|
if self.s is not None:
|
|
with suppress(Exception):
|
|
self.s.close()
|
|
|
|
def start(self):
|
|
"""
|
|
Starts the server to handle incoming connections.
|
|
"""
|
|
handler_thread = threading.Thread(target=self.elastic_client_handler)
|
|
handler_thread.daemon = True
|
|
handler_thread.start()
|
|
|
|
def elastic_client_handler(self):
|
|
"""
|
|
Handles incoming client connections.
|
|
"""
|
|
while True:
|
|
conn, addr = self.s.accept()
|
|
logger.info("Accept new connection from %s:%s...", *addr)
|
|
self.register_handler(conn, addr)
|
|
|
|
def register_handler(self, conn, addr, buffer_size=1024):
|
|
"""
|
|
Handles the registration of a client.
|
|
|
|
Parameters:
|
|
- conn: The connection socket.
|
|
- addr: The address of the client.
|
|
- buffer_size: The size of the buffer for receiving data.
|
|
"""
|
|
data_str = conn.recv(buffer_size).decode("utf-8")
|
|
if not data_str:
|
|
return
|
|
try:
|
|
data = json.loads(data_str)
|
|
except Exception:
|
|
logger.error("Failed to load %s as JSON string from %s", data_str, addr)
|
|
conn.close()
|
|
return
|
|
|
|
def is_valid_data(data):
|
|
"""
|
|
Validates the received data.
|
|
|
|
Parameters:
|
|
- data: The data to be validated.
|
|
|
|
Returns:
|
|
- True if the data is valid, otherwise False.
|
|
"""
|
|
if not isinstance(data, dict):
|
|
return False
|
|
if data.get("label") != "JOIN":
|
|
return False
|
|
content = data.get("content")
|
|
if not isinstance(content, dict):
|
|
return False
|
|
required_keys = ["device_id", "model_path", "tp", "pp", "port"]
|
|
if not all(k in content for k in required_keys):
|
|
return False
|
|
port = content["port"]
|
|
return isinstance(port, int) or (isinstance(port, str) and port.isdigit())
|
|
|
|
comm_name = None
|
|
if is_valid_data(data):
|
|
device_id = int(data["content"]["device_id"])
|
|
model_path = data["content"]["model_path"]
|
|
tp = int(data["content"]["tp"])
|
|
pp = int(data["content"]["pp"])
|
|
|
|
if (
|
|
int(self.device_id) == device_id
|
|
and self.model_path == model_path
|
|
and int(self.tp) == tp
|
|
and int(self.pp) == pp
|
|
):
|
|
comm_name = str(addr[0]) + ":" + str(addr[1])
|
|
ack = {"label": "JOIN_ACK", "content": {"name": comm_name}}
|
|
else:
|
|
server_desc = (int(self.device_id), self.model_path, int(self.tp), int(self.pp))
|
|
client_desc = (device_id, model_path, tp, pp)
|
|
msg = f"Received data {client_desc} does not consist with this server {server_desc}"
|
|
logger.warning(msg)
|
|
ack = {
|
|
"label": "JOIN_NACK",
|
|
"content": msg,
|
|
}
|
|
else:
|
|
logger.warning("Received data does not contain required fields: %s", data)
|
|
ack = {"label": "JOIN_NACK", "content": f"Received data does not contain required fields: {data}"}
|
|
|
|
try:
|
|
ack_str = json.dumps(ack).encode("utf-8")
|
|
except Exception as e:
|
|
logger.error("Failed to convert %s to JSON format, details: %s", ack, e)
|
|
conn.close()
|
|
return
|
|
|
|
try:
|
|
conn.send(ack_str)
|
|
except Exception as e:
|
|
logger.error("Failed to send %s to %s, details: %s", ack, addr, e)
|
|
conn.close()
|
|
return
|
|
|
|
if ack["content"] and isinstance(ack["content"], dict) and "name" in ack["content"]:
|
|
try:
|
|
p2psend = P2PSend(
|
|
self.addr,
|
|
data["content"]["port"],
|
|
ack["content"]["name"],
|
|
data["content"].get("group_name", "netloader"),
|
|
)
|
|
p2psend.send(self.model, self.original_int8)
|
|
except Exception as e:
|
|
logger.error("P2PSend Failed to send model to %s, details: %s", self.addr, e)
|
|
conn.close()
|