已合并
【FEAT】Simplify rank table parsing and add IPv6 support for HCCL collect #4989
刘家旺创建于 3月10日
【FEAT】Simplify rank table parsing and add IPv6 support for HCCL collect #4989
已合并
共 6 个文件变更+694-200
| @@ -21,8 +21,11 @@ import shutil | |||
| 21 | import subprocess | 21 | import subprocess |
| 22 | from abc import abstractmethod | 22 | from abc import abstractmethod |
| 23 | from concurrent.futures import ThreadPoolExecutor, as_completed | 23 | from concurrent.futures import ThreadPoolExecutor, as_completed |
| 24 | +import ipaddress | ||
| 25 | +from typing import Union | ||
| 24 | 26 | ||
| 25 | from msguard import Rule | 27 | from msguard import Rule |
| 28 | +from packaging.version import Version | ||
| 26 | 29 | ||
| 27 | from .base import BaseCollector | 30 | from .base import BaseCollector |
| 28 | from ..utils import get_npu_count, is_in_container | 31 | from ..utils import get_npu_count, is_in_container |
| @@ -104,10 +107,20 @@ class HCCLCollector(BaseCollector): | |||
| 104 | super().__init__() | 107 | super().__init__() |
| 105 | self.rank_table = rank_table | 108 | self.rank_table = rank_table |
| 106 | self.npu_count = npu_count if npu_count else get_npu_count() | 109 | self.npu_count = npu_count if npu_count else get_npu_count() |
| 107 | - self.option = "-hccs_ping" if getattr(rank_table, 'version', "1.0") == "1.2" else "-ping" | 110 | + self.option = ( |
| 111 | + "-hccs_ping" | ||
| 112 | + if getattr(rank_table, "version", Version("1.0")) >= Version("1.2") | ||
| 113 | + else "-ping" | ||
| 114 | + ) | ||
| 115 | + | ||
| 116 | + def _run_cmd(self, device_id: int, device_ip: Union[ipaddress.IPv4Address, ipaddress.IPv6Address]): | ||
| 117 | + if device_ip.version == 4: | ||
| 118 | + cmd = f"{HCCN_TOOL_CMD} -i {device_id} {self.option} -g address {device_ip}" | ||
| 119 | + elif device_ip.version == 6: | ||
| 120 | + cmd = f"{HCCN_TOOL_CMD} -i {device_id} {self.option} -inet6 -g ipv6_address {device_ip}" | ||
| 121 | + else: | ||
| 122 | + raise ValueError(f"Invalid IP version: {device_ip.version}") | ||
| 108 | 123 | ||
| 109 | - def _run_cmd(self, device_id: int, device_ip: str): | ||
| 110 | - cmd = f"{HCCN_TOOL_CMD} -i {device_id} {self.option} -g address {device_ip}" | ||
| 111 | proc = subprocess.Popen( | 124 | proc = subprocess.Popen( |
| 112 | shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True | 125 | shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True |
| 113 | ) | 126 | ) |
| @@ -68,8 +68,10 @@ from ..checkers import ( | |||
| 68 | from ..comparators import Comparator | 68 | from ..comparators import Comparator |
| 69 | from ..reporters import Reporter | 69 | from ..reporters import Reporter |
| 70 | from ..utils import ( | 70 | from ..utils import ( |
| 71 | - FrameworkType, ParserRegistry, update_model_type, | 71 | + Framework, update_model_type, |
| 72 | - CheckErrorHandler, ConfigErrorHandler, global_logger, singleton | 72 | + CheckErrorHandler, ConfigErrorHandler, global_logger, singleton, |
| 73 | + parse_rank_table | ||
| 74 | + | ||
| 73 | ) | 75 | ) |
| 74 | from ..cmate import inspect, run | 76 | from ..cmate import inspect, run |
| 75 | from ..cmate.cmate import _parse_configs, _parse_contexts | 77 | from ..cmate.cmate import _parse_configs, _parse_contexts |
| @@ -122,11 +124,10 @@ class CollectorFactory: | |||
| 122 | "msprechecker cannot determine the exact framework type of the rank table. " | 124 | "msprechecker cannot determine the exact framework type of the rank table. " |
| 123 | "Will use 'mindie' as the default framework." | 125 | "Will use 'mindie' as the default framework." |
| 124 | ) | 126 | ) |
| 125 | - args.framework = FrameworkType.TP_MINDIE | 127 | + args.framework = Framework.MINDIE.value |
| 126 | 128 | ||
| 127 | - framework_type = FrameworkType(args.framework) | 129 | + framework_type = Framework(args.framework) |
| 128 | - rank_table_parser = ParserRegistry.get(framework_type)() # create parser instance | 130 | + rank_table = parse_rank_table(args.rank_table_path, framework_type) |
| 129 | - rank_table = rank_table_parser.parse(args.rank_table_path) | ||
| 130 | 131 | ||
| 131 | collectors.extend( | 132 | collectors.extend( |
| 132 | ( | 133 | ( |
| @@ -17,7 +17,7 @@ | |||
| 17 | 17 | ||
| 18 | __all__ = [ | 18 | __all__ = [ |
| 19 | 'Version', 'get_pkg_version', | 19 | 'Version', 'get_pkg_version', |
| 20 | - 'get_npu_count', 'get_npu_type', 'RankTableParser', 'get_conn_mode', 'NpuType', | 20 | + 'get_npu_count', 'get_npu_type', 'get_conn_mode', 'NpuType', 'Framework', |
| 21 | 'get_current_ip_and_addr', | 21 | 'get_current_ip_and_addr', |
| 22 | "MacroExpander", "ExpandError", | 22 | "MacroExpander", "ExpandError", |
| 23 | 'Evaluator', | 23 | 'Evaluator', |
| @@ -25,16 +25,18 @@ __all__ = [ | |||
| 25 | 'is_in_container', 'singleton', | 25 | 'is_in_container', 'singleton', |
| 26 | 'global_logger', | 26 | 'global_logger', |
| 27 | 'Traverser', | 27 | 'Traverser', |
| 28 | - 'get_handler', | 28 | + 'get_handler', |
| 29 | - 'ErrorSeverity', 'ErrorType', | 29 | + 'ErrorSeverity', 'ErrorType', |
| 30 | - 'CollectError', 'BaseError', 'CheckError', | 30 | + 'CollectError', 'BaseError', 'CheckError', 'RankTableParseError', |
| 31 | - 'ErrorHandler', 'CollectErrorHandler', 'CheckErrorHandler', 'ConfigErrorHandler', 'CompareErrorHandler' | 31 | + 'ErrorHandler', 'CollectErrorHandler', 'CheckErrorHandler', 'ConfigErrorHandler', 'CompareErrorHandler', |
| 32 | + 'RankTable' | ||
| 32 | ] | 33 | ] |
| 33 | 34 | ||
| 34 | from .version import Version, get_pkg_version | 35 | from .version import Version, get_pkg_version |
| 35 | from .ascend import ( | 36 | from .ascend import ( |
| 36 | - get_npu_count, get_npu_type, get_conn_mode, NpuType, ParserRegistry, FrameworkType, | 37 | + get_npu_count, get_npu_type, get_conn_mode, NpuType, Framework, |
| 37 | - get_model_type, update_model_type | 38 | + get_model_type, update_model_type, parse_rank_table, RankTable, |
| 39 | + RankTableParseError | ||
| 38 | ) | 40 | ) |
| 39 | from .network import get_current_ip_and_addr | 41 | from .network import get_current_ip_and_addr |
| 40 | from .macro_expander import MacroExpander, ExpandError | 42 | from .macro_expander import MacroExpander, ExpandError |
| @@ -20,14 +20,17 @@ import os | |||
| 20 | import stat | 20 | import stat |
| 21 | import json | 21 | import json |
| 22 | import shlex | 22 | import shlex |
| 23 | +import ipaddress | ||
| 23 | import itertools | 24 | import itertools |
| 24 | import subprocess | 25 | import subprocess |
| 25 | from enum import Enum | 26 | from enum import Enum |
| 26 | -from abc import ABC, abstractmethod | 27 | +from pathlib import Path |
| 27 | from dataclasses import dataclass | 28 | from dataclasses import dataclass |
| 28 | -from typing import Dict, List, Type | 29 | +from typing import Dict, List, Union, Callable |
| 29 | 30 | ||
| 30 | from msguard.security import open_s | 31 | from msguard.security import open_s |
| 32 | +from packaging.version import Version, InvalidVersion | ||
| 33 | + | ||
| 31 | from .log import global_logger | 34 | from .log import global_logger |
| 32 | 35 | ||
| 33 | 36 | ||
| @@ -13,7 +13,7 @@ find = { exclude = ["test*"] } | |||
| 13 | # 运行相关 | 13 | # 运行相关 |
| 14 | [project] | 14 | [project] |
| 15 | name = "msprechecker" | 15 | name = "msprechecker" |
| 16 | -version = "0.3.4" | 16 | +version = "0.4.0" |
| 17 | description = "MindStudio Prechecker" | 17 | description = "MindStudio Prechecker" |
| 18 | authors = [ | 18 | authors = [ |
| 19 | {name = "AvadaKedavrua", email = "anonymousdev@163.com"}, | 19 | {name = "AvadaKedavrua", email = "anonymousdev@163.com"}, |
| @@ -35,7 +35,8 @@ dependencies = [ | |||
| 35 | "importlib_metadata; python_version == '3.7'", | 35 | "importlib_metadata; python_version == '3.7'", |
| 36 | "msguard", | 36 | "msguard", |
| 37 | "ply", | 37 | "ply", |
| 38 | - "colorama" | 38 | + "colorama", |
| 39 | + "packaging" | ||
| 39 | ] | 40 | ] |
| 40 | requires-python = ">=3.7" | 41 | requires-python = ">=3.7" |
| 41 | 42 | ||
| @@ -0,0 +1,386 @@ | |||
| 1 | +# ------------------------------------------------------------------------- | ||
| 2 | +# This file is part of the MindStudio project. | ||
| 3 | +# Copyright (c) 2025-2026 Huawei Technologies Co.,Ltd. | ||
| 4 | +# | ||
| 5 | +# MindStudio is licensed under Mulan PSL v2. | ||
| 6 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 7 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 8 | +# | ||
| 9 | +# `http://license.coscl.org.cn/MulanPSL2` | ||
| 10 | +# | ||
| 11 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 12 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 13 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 14 | +# See the Mulan PSL v2 for more details. | ||
| 15 | +# ------------------------------------------------------------------------- | ||
| 16 | + | ||
| 17 | +import json | ||
| 18 | +import pytest | ||
| 19 | +from unittest.mock import mock_open, patch | ||
| 20 | +from pathlib import Path | ||
| 21 | + | ||
| 22 | +from msprechecker.utils import parse_rank_table, Framework, RankTable, RankTableParseError | ||
| 23 | + | ||
| 24 | +# ============================================================================= | ||
| 25 | +# Tests for parse_rank_table | ||
| 26 | +# ============================================================================= | ||
| 27 | + | ||
| 28 | +@patch("builtins.open", new_callable=mock_open, read_data='{"server_list": [], "server_count": 0}') | ||
| 29 | +@patch("pathlib.Path.resolve") | ||
| 30 | +def test_parse_rank_table_given_mindie_framework_then_calls_mindie_parser( | ||
| 31 | + mock_resolve, mock_file | ||
| 32 | +): | ||
| 33 | + """Test parse_rank_table with MINDIE framework.""" | ||
| 34 | + mock_resolve.return_value = Path("/fake/path") | ||
| 35 | + with patch.object(Path, "is_file", return_value=True): | ||
| 36 | + # Test the framework dispatch - should not raise for empty server_list | ||
| 37 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 38 | + assert isinstance(result, RankTable) | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +@patch("builtins.open", new_callable=mock_open, read_data='{"prefill_device_list": [], "decode_device_list": []}') | ||
| 42 | +@patch("pathlib.Path.resolve") | ||
| 43 | +def test_parse_rank_table_given_vllm_framework_then_calls_vllm_parser( | ||
| 44 | + mock_resolve, mock_file | ||
| 45 | +): | ||
| 46 | + """Test parse_rank_table with VLLM framework.""" | ||
| 47 | + mock_resolve.return_value = Path("/fake/path") | ||
| 48 | + with patch.object(Path, "is_file", return_value=True): | ||
| 49 | + result = parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 50 | + assert isinstance(result, RankTable) | ||
| 51 | + | ||
| 52 | + | ||
| 53 | +def test_parse_rank_table_given_unknown_framework_then_raises_value_error(): | ||
| 54 | + """Test parse_rank_table with UNKNOWN framework.""" | ||
| 55 | + with pytest.raises(ValueError, match="No rank table parser"): | ||
| 56 | + parse_rank_table(Path("/fake/path"), Framework.UNKNOWN) | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +def test_parse_rank_table_given_sglang_framework_then_raises_value_error(): | ||
| 60 | + """Test parse_rank_table with SGLANG framework.""" | ||
| 61 | + with pytest.raises(ValueError, match="No rank table parser"): | ||
| 62 | + parse_rank_table(Path("/fake/path"), Framework.SGLANG) | ||
| 63 | + | ||
| 64 | + | ||
| 65 | + | ||
| 66 | +@patch("builtins.open", new_callable=mock_open) | ||
| 67 | +@patch("pathlib.Path.resolve") | ||
| 68 | +def test_parse_rank_table_mindie_with_valid_data(mock_resolve, mock_file): | ||
| 69 | + """Test parse_rank_table with MINDIE framework and valid data.""" | ||
| 70 | + mock_resolve.return_value = Path("/fake/path") | ||
| 71 | + data = { | ||
| 72 | + "server_list": [ | ||
| 73 | + { | ||
| 74 | + "server_id": "192.168.1.1", | ||
| 75 | + "device": [{"device_ip": "192.168.2.1", "device_id": 0, "rank_id": 0}], | ||
| 76 | + } | ||
| 77 | + ], | ||
| 78 | + "server_count": 1, | ||
| 79 | + "version": "1.0", | ||
| 80 | + } | ||
| 81 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 82 | + | ||
| 83 | + with patch.object(Path, "is_file", return_value=True): | ||
| 84 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 85 | + assert isinstance(result, RankTable) | ||
| 86 | + assert result.server_count == 1 | ||
| 87 | + | ||
| 88 | + | ||
| 89 | +@patch("builtins.open", new_callable=mock_open) | ||
| 90 | +@patch("pathlib.Path.resolve") | ||
| 91 | +def test_parse_rank_table_vllm_with_valid_data(mock_resolve, mock_file): | ||
| 92 | + """Test parse_rank_table with VLLM framework and valid data.""" | ||
| 93 | + mock_resolve.return_value = Path("/fake/path") | ||
| 94 | + data = { | ||
| 95 | + "prefill_device_list": [ | ||
| 96 | + { | ||
| 97 | + "server_id": "192.168.1.1", | ||
| 98 | + "device_ip": "192.168.2.1", | ||
| 99 | + "device_id": 0, | ||
| 100 | + "cluster_id": 1, | ||
| 101 | + } | ||
| 102 | + ], | ||
| 103 | + "decode_device_list": [], | ||
| 104 | + "server_count": 1, | ||
| 105 | + "version": "1.0", | ||
| 106 | + } | ||
| 107 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 108 | + | ||
| 109 | + with patch.object(Path, "is_file", return_value=True): | ||
| 110 | + result = parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 111 | + assert isinstance(result, RankTable) | ||
| 112 | + assert result.server_count == 1 | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +@patch("builtins.open", new_callable=mock_open) | ||
| 116 | +@patch("pathlib.Path.resolve") | ||
| 117 | +def test_parse_rank_table_mindie_with_invalid_server_id(mock_resolve, mock_file): | ||
| 118 | + mock_resolve.return_value = Path("/fake/path") | ||
| 119 | + data = { | ||
| 120 | + "server_list": [ | ||
| 121 | + { | ||
| 122 | + "server_id": "invalid_ip", | ||
| 123 | + "device": [{"device_ip": "192.168.2.1", "device_id": 0, "rank_id": 0}], | ||
| 124 | + } | ||
| 125 | + ], | ||
| 126 | + "server_count": 1, | ||
| 127 | + "version": "1.0", | ||
| 128 | + } | ||
| 129 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 130 | + | ||
| 131 | + with patch.object(Path, "is_file", return_value=True): | ||
| 132 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 133 | + assert isinstance(result, RankTable) | ||
| 134 | + | ||
| 135 | + | ||
| 136 | +@patch("builtins.open", new_callable=mock_open) | ||
| 137 | +@patch("pathlib.Path.resolve") | ||
| 138 | +def test_parse_rank_table_mindie_with_invalid_device_ip(mock_resolve, mock_file): | ||
| 139 | + mock_resolve.return_value = Path("/fake/path") | ||
| 140 | + data = { | ||
| 141 | + "server_list": [ | ||
| 142 | + { | ||
| 143 | + "server_id": "192.168.1.1", | ||
| 144 | + "device": [{"device_ip": "invalid_ip", "device_id": 0, "rank_id": 0}], | ||
| 145 | + } | ||
| 146 | + ], | ||
| 147 | + "server_count": 1, | ||
| 148 | + "version": "1.0", | ||
| 149 | + } | ||
| 150 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 151 | + | ||
| 152 | + with patch.object(Path, "is_file", return_value=True): | ||
| 153 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 154 | + assert isinstance(result, RankTable) | ||
| 155 | + | ||
| 156 | + | ||
| 157 | +@patch("builtins.open", new_callable=mock_open) | ||
| 158 | +@patch("pathlib.Path.resolve") | ||
| 159 | +def test_parse_rank_table_mindie_with_invalid_device_id(mock_resolve, mock_file): | ||
| 160 | + mock_resolve.return_value = Path("/fake/path") | ||
| 161 | + data = { | ||
| 162 | + "server_list": [ | ||
| 163 | + { | ||
| 164 | + "server_id": "192.168.1.1", | ||
| 165 | + "device": [ | ||
| 166 | + {"device_ip": "192.168.2.1", "device_id": "invalid", "rank_id": 0} | ||
| 167 | + ], | ||
| 168 | + } | ||
| 169 | + ], | ||
| 170 | + "server_count": 1, | ||
| 171 | + "version": "1.0", | ||
| 172 | + } | ||
| 173 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 174 | + | ||
| 175 | + with patch.object(Path, "is_file", return_value=True): | ||
| 176 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 177 | + assert isinstance(result, RankTable) | ||
| 178 | + | ||
| 179 | + | ||
| 180 | +@patch("builtins.open", new_callable=mock_open) | ||
| 181 | +@patch("pathlib.Path.resolve") | ||
| 182 | +def test_parse_rank_table_mindie_with_invalid_rank_id(mock_resolve, mock_file): | ||
| 183 | + mock_resolve.return_value = Path("/fake/path") | ||
| 184 | + data = { | ||
| 185 | + "server_list": [ | ||
| 186 | + { | ||
| 187 | + "server_id": "192.168.1.1", | ||
| 188 | + "device": [ | ||
| 189 | + {"device_ip": "192.168.2.1", "device_id": 0, "rank_id": "invalid"} | ||
| 190 | + ], | ||
| 191 | + } | ||
| 192 | + ], | ||
| 193 | + "server_count": 1, | ||
| 194 | + "version": "1.0", | ||
| 195 | + } | ||
| 196 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 197 | + | ||
| 198 | + with patch.object(Path, "is_file", return_value=True): | ||
| 199 | + result = parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 200 | + assert isinstance(result, RankTable) | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +@patch("builtins.open", new_callable=mock_open) | ||
| 204 | +@patch("pathlib.Path.resolve") | ||
| 205 | +def test_parse_rank_table_vllm_with_invalid_server_id(mock_resolve, mock_file): | ||
| 206 | + mock_resolve.return_value = Path("/fake/path") | ||
| 207 | + data = { | ||
| 208 | + "prefill_device_list": [ | ||
| 209 | + { | ||
| 210 | + "server_id": "invalid_ip", | ||
| 211 | + "device_ip": "192.168.2.1", | ||
| 212 | + "device_id": 0, | ||
| 213 | + "cluster_id": 1, | ||
| 214 | + } | ||
| 215 | + ], | ||
| 216 | + "decode_device_list": [], | ||
| 217 | + "server_count": 1, | ||
| 218 | + "version": "1.0", | ||
| 219 | + } | ||
| 220 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 221 | + | ||
| 222 | + with patch.object(Path, "is_file", return_value=True): | ||
| 223 | + result = parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 224 | + assert isinstance(result, RankTable) | ||
| 225 | + | ||
| 226 | + | ||
| 227 | +@patch("builtins.open", new_callable=mock_open) | ||
| 228 | +@patch("pathlib.Path.resolve") | ||
| 229 | +def test_parse_rank_table_vllm_with_invalid_cluster_id(mock_resolve, mock_file): | ||
| 230 | + mock_resolve.return_value = Path("/fake/path") | ||
| 231 | + data = { | ||
| 232 | + "prefill_device_list": [ | ||
| 233 | + { | ||
| 234 | + "server_id": "192.168.1.1", | ||
| 235 | + "device_ip": "192.168.2.1", | ||
| 236 | + "device_id": 0, | ||
| 237 | + "cluster_id": "invalid", | ||
| 238 | + } | ||
| 239 | + ], | ||
| 240 | + "decode_device_list": [], | ||
| 241 | + "server_count": 1, | ||
| 242 | + "version": "1.0", | ||
| 243 | + } | ||
| 244 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 245 | + | ||
| 246 | + with patch.object(Path, "is_file", return_value=True): | ||
| 247 | + result = parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 248 | + assert isinstance(result, RankTable) | ||
| 249 | + | ||
| 250 | + | ||
| 251 | +@patch("builtins.open", new_callable=mock_open) | ||
| 252 | +@patch("pathlib.Path.resolve") | ||
| 253 | +def test_parse_rank_table_mindie_exceeds_host_limit(mock_resolve, mock_file): | ||
| 254 | + """Test parse_rank_table with exceeds host limit.""" | ||
| 255 | + mock_resolve.return_value = Path("/fake/path") | ||
| 256 | + # Create data with more hosts than allowed | ||
| 257 | + data = { | ||
| 258 | + "server_list": [ | ||
| 259 | + {"server_id": f"192.168.{i}.1", "device": []} for i in range(1001) | ||
| 260 | + ], | ||
| 261 | + "server_count": 1001, | ||
| 262 | + "version": "1.0", | ||
| 263 | + } | ||
| 264 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 265 | + | ||
| 266 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 267 | + RankTableParseError, match="Host count exceeds limit" | ||
| 268 | + ): | ||
| 269 | + parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 270 | + | ||
| 271 | + | ||
| 272 | +@patch("builtins.open", new_callable=mock_open) | ||
| 273 | +@patch("pathlib.Path.resolve") | ||
| 274 | +def test_parse_rank_table_mindie_exceeds_device_limit(mock_resolve, mock_file): | ||
| 275 | + """Test parse_rank_table with exceeds device limit.""" | ||
| 276 | + mock_resolve.return_value = Path("/fake/path") | ||
| 277 | + # Create data with more devices than allowed per host | ||
| 278 | + data = { | ||
| 279 | + "server_list": [ | ||
| 280 | + { | ||
| 281 | + "server_id": "192.168.1.1", | ||
| 282 | + "device": [ | ||
| 283 | + {"device_ip": f"192.168.2.{i}", "device_id": i, "rank_id": i} | ||
| 284 | + for i in range(33) | ||
| 285 | + ], | ||
| 286 | + } | ||
| 287 | + ], | ||
| 288 | + "server_count": 1, | ||
| 289 | + "version": "1.0", | ||
| 290 | + } | ||
| 291 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 292 | + | ||
| 293 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 294 | + RankTableParseError, match="Device count for host .* exceeds limit" | ||
| 295 | + ): | ||
| 296 | + parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||
| 297 | + | ||
| 298 | + | ||
| 299 | +@patch("builtins.open", new_callable=mock_open) | ||
| 300 | +@patch("pathlib.Path.resolve") | ||
| 301 | +def test_parse_rank_table_vllm_exceeds_total_limit(mock_resolve, mock_file): | ||
| 302 | + """Test parse_rank_table with exceeds total limit.""" | ||
| 303 | + mock_resolve.return_value = Path("/fake/path") | ||
| 304 | + # Create data with more devices than total allowed | ||
| 305 | + data = { | ||
| 306 | + "prefill_device_list": [ | ||
| 307 | + { | ||
| 308 | + "server_id": f"192.168.{i}.1", | ||
| 309 | + "device_ip": f"192.168.{i}.2", | ||
| 310 | + "device_id": i, | ||
| 311 | + "cluster_id": i + 1, | ||
| 312 | + } | ||
| 313 | + for i in range(32001) # Exceeds _HOST_LIMIT * _DEVICE_LIMIT_PER_HOST | ||
| 314 | + ], | ||
| 315 | + "decode_device_list": [], | ||
| 316 | + "server_count": 1, | ||
| 317 | + "version": "1.0", | ||
| 318 | + } | ||
| 319 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 320 | + | ||
| 321 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 322 | + RankTableParseError, match="length exceeds limit" | ||
| 323 | + ): | ||
| 324 | + parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 325 | + | ||
| 326 | + | ||
| 327 | +@patch("builtins.open", new_callable=mock_open) | ||
| 328 | +@patch("pathlib.Path.resolve") | ||
| 329 | +def test_parse_rank_table_vllm_exceeds_host_limit(mock_resolve, mock_file): | ||
| 330 | + """Test parse_rank_table with exceeds host limit.""" | ||
| 331 | + mock_resolve.return_value = Path("/fake/path") | ||
| 332 | + # Create data with more hosts than allowed (using valid IPs) | ||
| 333 | + devices = [] | ||
| 334 | + for i in range(1001): | ||
| 335 | + # Use valid IPs in the 10.0.x.x range | ||
| 336 | + octet2 = i // 256 | ||
| 337 | + octet3 = i % 256 | ||
| 338 | + devices.append( | ||
| 339 | + { | ||
| 340 | + "server_id": f"10.0.{octet2}.{octet3}", | ||
| 341 | + "device_ip": f"10.1.{octet2}.{octet3}", | ||
| 342 | + "device_id": 0, | ||
| 343 | + "cluster_id": i + 1, | ||
| 344 | + } | ||
| 345 | + ) | ||
| 346 | + | ||
| 347 | + data = { | ||
| 348 | + "prefill_device_list": devices, | ||
| 349 | + "decode_device_list": [], | ||
| 350 | + "server_count": 1001, | ||
| 351 | + "version": "1.0", | ||
| 352 | + } | ||
| 353 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 354 | + | ||
| 355 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 356 | + RankTableParseError, match=r"Host count exceeds limit \d+" | ||
| 357 | + ): | ||
| 358 | + parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 359 | + | ||
| 360 | + | ||
| 361 | +@patch("builtins.open", new_callable=mock_open) | ||
| 362 | +@patch("pathlib.Path.resolve") | ||
| 363 | +def test_parse_rank_table_vllm_missing_device_list(mock_resolve, mock_file): | ||
| 364 | + """Test parse_rank_table with missing device list.""" | ||
| 365 | + mock_resolve.return_value = Path("/fake/path") | ||
| 366 | + data = { | ||
| 367 | + "server_count": 1, | ||
| 368 | + "version": "1.0", | ||
| 369 | + } | ||
| 370 | + mock_file.return_value.read.return_value = json.dumps(data) | ||
| 371 | + | ||
| 372 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 373 | + RankTableParseError, match="Expected 'prefill_device_list' and 'decode_device_list'" | ||
| 374 | + ): | ||
| 375 | + parse_rank_table(Path("/fake/path"), Framework.VLLM) | ||
| 376 | + | ||
| 377 | + | ||
| 378 | +@patch("builtins.open", side_effect=OSError("Failed to open")) | ||
| 379 | +@patch("pathlib.Path.resolve") | ||
| 380 | +def test_parse_rank_table_json_load_error(mock_resolve, mock_file): | ||
| 381 | + """Test parse_rank_table with file open error.""" | ||
| 382 | + mock_resolve.return_value = Path("/fake/path") | ||
| 383 | + with patch.object(Path, "is_file", return_value=True), pytest.raises( | ||
| 384 | + RankTableParseError | ||
| 385 | + ): | ||
| 386 | + parse_rank_table(Path("/fake/path"), Framework.MINDIE) | ||