438 lines
15 KiB
Python
438 lines
15 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 io
|
|
import json
|
|
import logging
|
|
import socket
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_ascend.model_loader.netloader.interaction import elastic
|
|
from vllm_ascend.model_loader.netloader.interaction.elastic import ElasticClient, ElasticServer
|
|
|
|
|
|
# Simulate server's normal response
|
|
def mock_server_response(data):
|
|
return json.dumps({"label": "JOIN_ACK", "content": {"name": "mocked_name"}}).encode("utf-8")
|
|
|
|
|
|
# Simulate server's error response
|
|
def mock_server_error_response(data):
|
|
return json.dumps({"label": "JOIN_ACK", "content": None}).encode("utf-8")
|
|
|
|
|
|
# Simulated server's abnormal response
|
|
def mock_server_exception_response(data):
|
|
raise Exception("Mocked server exception")
|
|
|
|
|
|
@contextmanager
|
|
def capture_elastic_logs(level=logging.DEBUG):
|
|
log_capture_string = io.StringIO()
|
|
handler = logging.StreamHandler(log_capture_string)
|
|
handler.setLevel(level)
|
|
original_level = elastic.logger.level
|
|
elastic.logger.setLevel(level)
|
|
elastic.logger.addHandler(handler)
|
|
try:
|
|
yield log_capture_string
|
|
finally:
|
|
elastic.logger.removeHandler(handler)
|
|
elastic.logger.setLevel(original_level)
|
|
log_capture_string.close()
|
|
|
|
|
|
# Test the initialization of ElasticClient
|
|
def test_elastic_client_init():
|
|
sources = ["127.0.0.1:12345"]
|
|
device_id = 0
|
|
model_path = "mocked_model_path"
|
|
tp = 1
|
|
pp = 1
|
|
|
|
with patch("socket.socket") as mock_socket:
|
|
mock_socket_instance = MagicMock()
|
|
mock_socket.return_value = mock_socket_instance
|
|
mock_socket_instance.recv.return_value = mock_server_response(None)
|
|
|
|
mock_socket_instance.getsockname.return_value = ("127.0.0.1", 12346)
|
|
mock_socket_instance.__enter__.return_value = mock_socket_instance
|
|
|
|
with ElasticClient(sources, device_id, model_path, tp, pp) as client:
|
|
assert client.server_addr == "127.0.0.1"
|
|
assert client.server_port == 12345
|
|
assert client.ack == ("mocked_name", 12346)
|
|
mock_socket_instance.close.assert_called_once()
|
|
|
|
|
|
# Test the register method of ElasticClient
|
|
def test_elastic_client_register():
|
|
sources = ["127.0.0.1:12345"]
|
|
device_id = 0
|
|
model_path = "mocked_model_path"
|
|
tp = 1
|
|
pp = 1
|
|
|
|
with patch("socket.socket") as mock_socket:
|
|
mock_socket_instance = MagicMock()
|
|
mock_socket.return_value = mock_socket_instance
|
|
mock_socket_instance.connect.return_value = None
|
|
mock_socket_instance.recv.return_value = mock_server_response(None)
|
|
|
|
mock_socket_instance.getsockname.return_value = ("127.0.0.1", 12346)
|
|
mock_socket_instance.__enter__.return_value = mock_socket_instance
|
|
|
|
client = ElasticClient(sources, device_id, model_path, tp, pp)
|
|
assert client.register(device_id, model_path, tp, pp) == ("mocked_name", 12346)
|
|
|
|
|
|
# Test the behavior of the `register` method of ElasticClient when the server returns an error response.
|
|
def test_elastic_client_register_error_response():
|
|
sources = ["127.0.0.1:12345"]
|
|
device_id = 0
|
|
model_path = "mocked_model_path"
|
|
tp = 1
|
|
pp = 1
|
|
|
|
with patch("socket.socket") as mock_socket:
|
|
mock_socket_instance = MagicMock()
|
|
mock_socket.return_value = mock_socket_instance
|
|
mock_socket_instance.connect.return_value = None
|
|
mock_socket_instance.recv.return_value = mock_server_error_response(None)
|
|
|
|
with ElasticClient(sources, device_id, model_path, tp, pp) as client, pytest.raises(RuntimeError):
|
|
client.register(device_id, model_path, tp, pp)
|
|
mock_socket_instance.close.assert_called_once()
|
|
|
|
|
|
# Test the behavior of the `register` method of ElasticClient when an exception is thrown on the server.
|
|
def test_elastic_client_register_exception():
|
|
sources = ["127.0.0.1:12345"]
|
|
device_id = 0
|
|
model_path = "mocked_model_path"
|
|
tp = 1
|
|
pp = 1
|
|
|
|
with patch("socket.socket") as mock_socket:
|
|
mock_socket_instance = MagicMock()
|
|
mock_socket.return_value = mock_socket_instance
|
|
mock_socket_instance.connect.return_value = None
|
|
mock_socket_instance.recv.side_effect = mock_server_exception_response
|
|
mock_socket_instance.__enter__.return_value = mock_socket_instance
|
|
mock_socket_instance.__exit__.return_value = None
|
|
|
|
with ElasticClient(sources, device_id, model_path, tp, pp) as client, pytest.raises(RuntimeError):
|
|
client.register(device_id, model_path, tp, pp)
|
|
mock_socket_instance.close.assert_called_once()
|
|
|
|
|
|
class FakeInt8Param:
|
|
def __init__(self, name="param", device="npu", dtype=torch.int8):
|
|
self.dtype = dtype
|
|
self.device = torch.device(device)
|
|
|
|
@property
|
|
def data(self):
|
|
return self # Simulate .data returning self so .cpu() etc. can be chained
|
|
|
|
def clone(self):
|
|
return self
|
|
|
|
def detach(self):
|
|
return self
|
|
|
|
def cpu(self):
|
|
self.device = torch.device("cpu")
|
|
return self
|
|
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.params = {
|
|
"param1": MagicMock(dtype=torch.float32), # This will be ignored
|
|
"param2": FakeInt8Param(), # This simulates a real int8 param
|
|
}
|
|
|
|
def named_parameters(self):
|
|
return self.params.items()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_model():
|
|
return FakeModel()
|
|
|
|
|
|
@pytest.fixture
|
|
def server_config():
|
|
return {
|
|
"addr": "127.0.0.1",
|
|
"port": 8080,
|
|
"model": MagicMock(),
|
|
"device_id": 0,
|
|
"model_path": "/test/model",
|
|
"tp": 1,
|
|
"pp": 1,
|
|
"int8_cache": "dram",
|
|
"int8_cache_name": None,
|
|
}
|
|
|
|
|
|
# Test server initialization
|
|
def test_server_initialization(server_config, mock_model):
|
|
server_config["model"] = mock_model
|
|
with patch("socket.socket") as mock_socket, capture_elastic_logs() as log_capture_string:
|
|
server = ElasticServer(**server_config)
|
|
|
|
# Check the socket configuration
|
|
mock_socket.assert_called_with(socket.AF_INET, socket.SOCK_STREAM)
|
|
mock_socket.return_value.setsockopt.assert_called_with(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
mock_socket.return_value.bind.assert_called_with(("127.0.0.1", 8080))
|
|
mock_socket.return_value.listen.assert_called_with(256)
|
|
|
|
# Check int8 cache
|
|
assert "param2" in server.original_int8
|
|
assert server.original_int8["param2"].device.type == "cpu" # Verifying DRAM Cache
|
|
|
|
assert server.addr == server_config["addr"]
|
|
assert server.port == server_config["port"]
|
|
assert server.device_id == server_config["device_id"]
|
|
assert server.model_path == server_config["model_path"]
|
|
assert server.tp == server_config["tp"]
|
|
assert server.pp == server_config["pp"]
|
|
|
|
# Get captured logs
|
|
log_output = log_capture_string.getvalue()
|
|
|
|
# Check output
|
|
assert "Server 127.0.0.1:8080 starts" in log_output
|
|
|
|
|
|
# Test the int8 cache option
|
|
@pytest.mark.parametrize("cache_option,expected_device", [("dram", "cpu"), ("no", None), ("invalid", None)])
|
|
def test_int8_cache_handling(server_config, mock_model, cache_option, expected_device):
|
|
server_config["int8_cache"] = cache_option
|
|
server_config["model"] = mock_model
|
|
|
|
with patch("socket.socket"), capture_elastic_logs() as log_capture_string:
|
|
server = ElasticServer(**server_config)
|
|
|
|
log_output = log_capture_string.getvalue()
|
|
|
|
if cache_option == "invalid":
|
|
assert "int8_cache should be selected in [HBM, DRAM]" in log_output
|
|
|
|
if expected_device is None:
|
|
assert len(server.original_int8) == 0
|
|
else:
|
|
assert server.original_int8["param2"].device.type == expected_device
|
|
|
|
|
|
# Test client processing
|
|
def test_client_handler_valid_join(server_config, mock_model):
|
|
server_config["model"] = mock_model
|
|
with patch("vllm_ascend.model_loader.netloader.interaction.elastic.P2PSend") as mock_p2p_send:
|
|
# Create a simulated connection
|
|
mock_conn = MagicMock()
|
|
mock_addr = ("192.168.1.1", 12345)
|
|
|
|
# Configuring Client Data
|
|
valid_data = {
|
|
"label": "JOIN",
|
|
"content": {"device_id": 0, "model_path": "/test/model", "tp": 1, "pp": 1, "port": 9090},
|
|
}
|
|
mock_conn.recv.return_value = json.dumps(valid_data).encode("utf-8")
|
|
|
|
# Start the server
|
|
server = ElasticServer(**server_config)
|
|
server.register_handler(mock_conn, mock_addr)
|
|
|
|
# Verify response
|
|
expected_ack = {"label": "JOIN_ACK", "content": {"name": "192.168.1.1:12345"}}
|
|
mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
|
|
mock_p2p_send.assert_called_once_with("127.0.0.1", 9090, "192.168.1.1:12345", "netloader")
|
|
mock_conn.close.assert_called_once()
|
|
|
|
|
|
# Test mismatched JOIN requests
|
|
def test_client_handler_mismatch(server_config):
|
|
with patch("socket.socket"):
|
|
server = ElasticServer(**server_config)
|
|
mock_conn = MagicMock()
|
|
mock_addr = ("192.168.1.1", 12345)
|
|
|
|
# Send mismatched data
|
|
mismatch_data = {
|
|
"label": "JOIN",
|
|
"content": {
|
|
"device_id": 1, # 不匹配的ID
|
|
"model_path": "/wrong/model",
|
|
"tp": 2,
|
|
"pp": 2,
|
|
"port": 9090,
|
|
},
|
|
}
|
|
mock_conn.recv.return_value = json.dumps(mismatch_data).encode("utf-8")
|
|
|
|
server.register_handler(mock_conn, mock_addr)
|
|
|
|
assert isinstance(mismatch_data["content"], dict)
|
|
|
|
# Verify response
|
|
mismatch_tuple = (
|
|
mismatch_data["content"]["device_id"],
|
|
mismatch_data["content"]["model_path"],
|
|
mismatch_data["content"]["tp"],
|
|
mismatch_data["content"]["pp"],
|
|
)
|
|
|
|
server_tuple = (
|
|
server_config["device_id"],
|
|
server_config["model_path"],
|
|
server_config["tp"],
|
|
server_config["pp"],
|
|
)
|
|
|
|
expected_ack = {
|
|
"label": "JOIN_NACK",
|
|
"content": (f"Received data {mismatch_tuple} does not consist with this server {server_tuple}"),
|
|
}
|
|
mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
|
|
mock_conn.close.assert_called_once()
|
|
|
|
|
|
# Test Invalid Request
|
|
@pytest.mark.parametrize(
|
|
"invalid_data,should_send",
|
|
[
|
|
({"label": "WRONG_LABEL"}, True), # Incorrect label, can be decoded as JSON, but the content is invalid.
|
|
(
|
|
{"content": {"missing_fields": True}},
|
|
True,
|
|
), # Missing field, can be decoded as JSON, but the content is invalid.
|
|
("plain text", False), # Non-JSON data, json.loads failed
|
|
(b"invalid_bytes", False), # Invalid byte, decode or json.loads failed
|
|
],
|
|
)
|
|
def test_client_handler_invalid_requests(server_config, invalid_data, should_send):
|
|
with patch("socket.socket"), capture_elastic_logs() as log_capture_string:
|
|
server = ElasticServer(**server_config)
|
|
mock_conn = MagicMock()
|
|
mock_addr = ("192.168.1.1", 12345)
|
|
|
|
if isinstance(invalid_data, (str, bytes)):
|
|
mock_conn.recv.return_value = invalid_data if isinstance(invalid_data, bytes) else invalid_data.encode()
|
|
else:
|
|
mock_conn.recv.return_value = json.dumps(invalid_data).encode("utf-8")
|
|
|
|
server.register_handler(mock_conn, mock_addr)
|
|
|
|
if should_send:
|
|
expected_ack = {
|
|
"label": "JOIN_NACK",
|
|
"content": f"Received data does not contain required fields: {invalid_data}",
|
|
}
|
|
mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
|
|
else:
|
|
mock_conn.send.assert_not_called()
|
|
|
|
log_output = log_capture_string.getvalue()
|
|
|
|
# Any warning in the log is acceptable
|
|
assert "Failed to load" in log_output or "does not contain" in log_output
|
|
mock_conn.close.assert_called_once()
|
|
|
|
|
|
# Test the thread startup.
|
|
def test_server_start(server_config):
|
|
with patch("socket.socket"), patch("threading.Thread") as mock_thread:
|
|
handler_thread_instance = mock_thread.return_value
|
|
|
|
server = ElasticServer(**server_config)
|
|
server.start()
|
|
|
|
# Assert that the correct target parameter was passed when instantiating the Thread instance.
|
|
mock_thread.assert_called_once()
|
|
args, kwargs = mock_thread.call_args
|
|
assert kwargs["target"] == server.elastic_client_handler
|
|
|
|
# Verify the daemon attribute is set to True (the attribute value will be recorded after MagicMock assignment).
|
|
assert handler_thread_instance.daemon is True
|
|
|
|
# Check if the start() method is called.
|
|
handler_thread_instance.start.assert_called_once()
|
|
|
|
|
|
# Test resource clearing
|
|
def test_server_cleanup(server_config):
|
|
with patch("socket.socket") as mock_socket:
|
|
server = ElasticServer(**server_config)
|
|
del server
|
|
mock_socket.return_value.close.assert_called_once()
|
|
|
|
|
|
def test_draft_group_name_in_client_register():
|
|
sent_payloads = []
|
|
|
|
with (
|
|
patch("socket.socket") as mock_socket,
|
|
patch("vllm_ascend.model_loader.netloader.interaction.elastic.find_free_port", return_value=12346),
|
|
):
|
|
mock_socket_instance = MagicMock()
|
|
mock_socket.return_value = mock_socket_instance
|
|
mock_socket_instance.recv.return_value = mock_server_response(None)
|
|
mock_socket_instance.send.side_effect = lambda data: sent_payloads.append(json.loads(data.decode()))
|
|
|
|
client = ElasticClient(["127.0.0.1:12345"], 0, "draft-model", 1, 1, "netloader_draft")
|
|
client.s = mock_socket_instance
|
|
client.register(0, "draft-model", 1, 1)
|
|
|
|
assert client.group_name == "netloader_draft"
|
|
assert sent_payloads[0]["content"]["group_name"] == "netloader_draft"
|
|
|
|
|
|
def test_draft_group_name_in_server_p2p_send(server_config, mock_model):
|
|
server_config["model"] = mock_model
|
|
join_data = {
|
|
"label": "JOIN",
|
|
"content": {
|
|
"device_id": 0,
|
|
"model_path": "/test/model",
|
|
"tp": 1,
|
|
"pp": 1,
|
|
"port": 9090,
|
|
"group_name": "netloader_draft",
|
|
},
|
|
}
|
|
|
|
with (
|
|
patch("socket.socket"),
|
|
patch("vllm_ascend.model_loader.netloader.interaction.elastic.P2PSend") as mock_p2p_send,
|
|
):
|
|
mock_conn = MagicMock()
|
|
mock_conn.recv.return_value = json.dumps(join_data).encode("utf-8")
|
|
|
|
ElasticServer(**server_config, group_name="netloader_draft").register_handler(mock_conn, ("192.168.1.1", 12345))
|
|
|
|
mock_p2p_send.assert_called_once_with("127.0.0.1", 9090, "192.168.1.1:12345", "netloader_draft")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|