mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-12-24 13:28:13 +08:00
0c8c6369ed
* [feat] simplify configuration for pd-disaggregated deployment, and refactor post-init and usage for all ports * [fix] fix some bugs * [fix] fix rdma port for cache manager/messager * [fix] temporarily cancel port availability check to see if it can pass ci test * [feat] simplify args for multi api server * [fix] fix dp * [fix] fix port for xpu * [fix] add tests for ports post processing & fix ci * [test] fix test_multi_api_server * [fix] fix rdma_comm_ports args for multi_api_server * [fix] fix test_common_engine * [fix] fix test_cache_transfer_manager * [chore] automatically setting FD_ENABLE_MULTI_API_SERVER * [fix] avoid api server from creating engine_args twice * [fix] fix test_run_batch * [fix] fix test_metrics * [fix] fix splitwise connector init * [test] add test_rdma_transfer and test_expert_service * [fix] fix code syntax * [fix] fix test_rdma_transfer and build wheel with rdma script
139 lines
6.0 KiB
Python
139 lines
6.0 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. 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 os
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
from fastdeploy.cache_manager.transfer_factory.rdma_cache_transfer import (
|
|
RDMACommManager,
|
|
)
|
|
|
|
|
|
class TestRDMACommManager(unittest.TestCase):
|
|
def setUp(self):
|
|
self.args = {
|
|
"splitwise_role": "prefill",
|
|
"rank": 0,
|
|
"gpu_id": 0,
|
|
"cache_k_ptr_list": [1, 2, 3],
|
|
"cache_v_ptr_list": [4, 5, 6],
|
|
"max_block_num": 10,
|
|
"block_bytes": 1024,
|
|
"rdma_port": 12345,
|
|
"prefill_tp_size": 1,
|
|
"prefill_tp_idx": 0,
|
|
}
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "", "KVCACHE_RDMA_NICS": ""})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_init_rdma_comm_manager_on_gpu_init_all(self, mock_run, mock_rdma_comm, mock_platform):
|
|
# Case: Automatically set all environment variables
|
|
mock_platform.is_cuda.return_value = True
|
|
mock_platform.device_name = "gpu"
|
|
mock_run.side_effect = [
|
|
Mock(returncode=0, stdout="8.0\n"),
|
|
Mock(returncode=0, stdout="KVCACHE_RDMA_NICS=mlx5_2\n"),
|
|
]
|
|
|
|
manager = RDMACommManager(**self.args)
|
|
self.assertEqual(manager.splitwise_role, "prefill")
|
|
self.assertEqual(mock_run.call_count, 2)
|
|
mock_rdma_comm.assert_called_once()
|
|
self.assertEqual(os.getenv("KVCACHE_GDRCOPY_FLUSH_ENABLE"), "1")
|
|
self.assertEqual(os.getenv("KVCACHE_RDMA_NICS"), "mlx5_2")
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "", "KVCACHE_RDMA_NICS": "mlx5_1"})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_init_rdma_comm_manager_on_gpu_init_gdrcopy(self, mock_run, mock_rdma_comm, mock_platform):
|
|
# Case: Only set KVCACHE_GDRCOPY_FLUSH_ENABLE
|
|
mock_platform.is_cuda.return_value = True
|
|
mock_platform.device_name = "gpu"
|
|
mock_run.side_effect = [Mock(returncode=0, stdout="8.0\n")]
|
|
|
|
manager = RDMACommManager(**self.args)
|
|
self.assertEqual(manager.splitwise_role, "prefill")
|
|
self.assertEqual(mock_run.call_count, 1)
|
|
mock_rdma_comm.assert_called_once()
|
|
self.assertEqual(os.getenv("KVCACHE_GDRCOPY_FLUSH_ENABLE"), "1")
|
|
self.assertEqual(os.getenv("KVCACHE_RDMA_NICS"), "mlx5_1")
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "0", "KVCACHE_RDMA_NICS": ""})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_init_rdma_comm_manager_on_gpu_init_nics(self, mock_run, mock_rdma_comm, mock_platform):
|
|
# Case: Only set KVCACHE_RDMA_NICS
|
|
mock_platform.is_cuda.return_value = True
|
|
mock_platform.device_name = "gpu"
|
|
mock_run.side_effect = [Mock(returncode=0, stdout="KVCACHE_RDMA_NICS=mlx5_2\n")]
|
|
|
|
manager = RDMACommManager(**self.args)
|
|
self.assertEqual(manager.splitwise_role, "prefill")
|
|
self.assertEqual(mock_run.call_count, 1)
|
|
mock_rdma_comm.assert_called_once()
|
|
self.assertEqual(os.getenv("KVCACHE_GDRCOPY_FLUSH_ENABLE"), "0")
|
|
self.assertEqual(os.getenv("KVCACHE_RDMA_NICS"), "mlx5_2")
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "0", "KVCACHE_RDMA_NICS": "mlx5_1"})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_init_rdma_comm_manager_on_gpu_init_nothing(self, mock_run, mock_rdma_comm, mock_platform):
|
|
# Case: Do not set any environment variables
|
|
mock_platform.is_cuda.return_value = True
|
|
mock_platform.device_name = "gpu"
|
|
|
|
manager = RDMACommManager(**self.args)
|
|
self.assertEqual(manager.splitwise_role, "prefill")
|
|
self.assertEqual(mock_run.call_count, 0)
|
|
mock_rdma_comm.assert_called_once()
|
|
self.assertEqual(os.getenv("KVCACHE_GDRCOPY_FLUSH_ENABLE"), "0")
|
|
self.assertEqual(os.getenv("KVCACHE_RDMA_NICS"), "mlx5_1")
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "0", "KVCACHE_RDMA_NICS": "mlx5_1"})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_connect_success(self, mock_run, mock_rdma_comm, mock_platform):
|
|
"""Test successful connection"""
|
|
manager = RDMACommManager(**self.args)
|
|
manager.messager.is_connected.return_value = False
|
|
manager.messager.connect.return_value = 0
|
|
|
|
result = manager.connect("127.0.0.1", 12345)
|
|
self.assertTrue(result)
|
|
manager.messager.connect.assert_called_once_with("127.0.0.1", "12345", 0)
|
|
|
|
@patch.dict("os.environ", {"KVCACHE_GDRCOPY_FLUSH_ENABLE": "0", "KVCACHE_RDMA_NICS": "mlx5_1"})
|
|
@patch("fastdeploy.platforms.current_platform")
|
|
@patch("rdma_comm.RDMACommunicator")
|
|
@patch("subprocess.run")
|
|
def test_write_cache(self, mock_run, mock_rdma_comm, mock_platform):
|
|
"""Test write_cache method"""
|
|
manager = RDMACommManager(**self.args)
|
|
manager.messager.write_cache.return_value = True
|
|
|
|
result = manager.write_cache("127.0.0.1", 12345, [1, 2], [3, 4], 0)
|
|
self.assertTrue(result)
|
|
manager.messager.write_cache.assert_called_once_with("127.0.0.1", "12345", [1, 2], [3, 4], 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|