已合并
【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
已合并
刘家旺创建于 3月10日
共 6 个文件变更+694-200
@@ -21,8 +21,11 @@ import shutil
21import subprocess21import subprocess
22from abc import abstractmethod22from abc import abstractmethod
23from concurrent.futures import ThreadPoolExecutor, as_completed23from concurrent.futures import ThreadPoolExecutor, as_completed
24+import ipaddress
25+from typing import Union
24 26 
25from msguard import Rule27from msguard import Rule
28+from packaging.version import Version
26 29 
27from .base import BaseCollector30from .base import BaseCollector
28from ..utils import get_npu_count, is_in_container31from ..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_table108 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=True125 shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
113 )126 )
@@ -68,8 +68,10 @@ from ..checkers import (
68from ..comparators import Comparator68from ..comparators import Comparator
69from ..reporters import Reporter69from ..reporters import Reporter
70from ..utils import (70from ..utils import (
71- FrameworkType, ParserRegistry, update_model_type,71+ Framework, update_model_type,
72- CheckErrorHandler, ConfigErrorHandler, global_logger, singleton72+ CheckErrorHandler, ConfigErrorHandler, global_logger, singleton,
73+ parse_rank_table
74+
73)75)
74from ..cmate import inspect, run76from ..cmate import inspect, run
75from ..cmate.cmate import _parse_configs, _parse_contexts77from ..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_MINDIE127+ 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 instance130+ 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 
34from .version import Version, get_pkg_version35from .version import Version, get_pkg_version
35from .ascend import (36from .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_type38+ get_model_type, update_model_type, parse_rank_table, RankTable,
39+ RankTableParseError
38)40)
39from .network import get_current_ip_and_addr41from .network import get_current_ip_and_addr
40from .macro_expander import MacroExpander, ExpandError42from .macro_expander import MacroExpander, ExpandError
@@ -20,14 +20,17 @@ import os
20import stat20import stat
21import json21import json
22import shlex22import shlex
23+import ipaddress
23import itertools24import itertools
24import subprocess25import subprocess
25from enum import Enum26from enum import Enum
26-from abc import ABC, abstractmethod27+from pathlib import Path
27from dataclasses import dataclass28from dataclasses import dataclass
28-from typing import Dict, List, Type29+from typing import Dict, List, Union, Callable
29 30 
30from msguard.security import open_s31from msguard.security import open_s
32+from packaging.version import Version, InvalidVersion
33+ 
31from .log import global_logger34from .log import global_logger
32 35 
33 36 
@@ -13,7 +13,7 @@ find = { exclude = ["test*"] }
13# 运行相关13# 运行相关
14[project]14[project]
15name = "msprechecker"15name = "msprechecker"
16-version = "0.3.4"16+version = "0.4.0"
17description = "MindStudio Prechecker"17description = "MindStudio Prechecker"
18authors = [18authors = [
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]
40requires-python = ">=3.7"41requires-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)